Merge pull request #1266 from mito-ds/parser-index-labels-bug

Fix bug in parsing specific index labels

Author: marthacryan

mitosheet/mitosheet/parser.py

@@ -1045,12 +1045,7 @@ def parse_formula(
             if is_datetime_index(df.index) and is_prev_version(get_pandas_version(), '1.0.0'):
                 index_labels = pd.to_datetime(index_labels)
 
-            if len(column_header_dependencies) > 0:
-                final_set_code = f'({code_with_functions}).loc[{get_column_header_list_as_transpiled_code(index_labels)}]' # type: ignore
-            else:
-                final_set_code = f'{code_with_functions}'
-                
-            final_code = f'{df_name}.loc[{get_column_header_list_as_transpiled_code(index_labels)}, [{transpiled_column_header}]] = {final_set_code}' # type: ignore
+            final_code = f'{df_name}.loc[{get_column_header_list_as_transpiled_code(index_labels)}, [{transpiled_column_header}]] = {code_with_functions}' # type: ignore
 
     else:
         final_code = f'{code_with_functions}'

mitosheet/mitosheet/tests/step_performers/column_steps/test_set_column_formula.py

@@ -535,6 +535,20 @@ def test_set_specific_index_labels_twice():
 
     assert mito.dfs[0].equals(pd.DataFrame({'A': [1, 2, 3], 'B': [2, 0, 0]}))
 
+def test_set_specific_index_labels_to_header_header():
+    mito = create_mito_wrapper(pd.DataFrame({'A': [1, 2, 3]}))
+    mito.add_column(0, 'B')
+    mito.set_formula('=SUM(A:A)', 0, 'B', index_labels=[0])
+
+    assert mito.dfs[0].equals(pd.DataFrame({'A': [1, 2, 3], 'B': [6, 0, 0]}))
+
+
+def test_set_specific_index_labels_to_header_header_multiple_header_dependencies():
+    mito = create_mito_wrapper(pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]}))
+    mito.add_column(0, 'C')
+    mito.set_formula('=SUM(A:B)', 0, 'C', index_labels=[0])
+
+    assert mito.dfs[0].equals(pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6], 'C': [21, 0, 0]}))
 
 CROSS_SHEET_TESTS = [
     (

mitosheet/mitosheet/tests/test_parse.py

@@ -11,7 +11,7 @@
 
 from mitosheet.errors import MitoError
 from mitosheet.parser import get_backend_formula_from_frontend_formula, parse_formula, safe_contains, get_frontend_formula
-from mitosheet.types import FORMULA_ENTIRE_COLUMN_TYPE
+from mitosheet.types import FORMULA_ENTIRE_COLUMN_TYPE, FORMULA_SPECIFIC_INDEX_LABELS_TYPE
 from mitosheet.tests.decorators import pandas_post_1_2_only
 
 
@@ -1274,6 +1274,38 @@ def test_parse_cross_sheet_formulas(formula, column_header, formula_label, dfs,
     )
 ]
 
+def test_specific_index_labels_header_header():
+    formula = '=SUM(A:A)'
+    column_header = 'B'
+    formula_label = 0
+    df = pd.DataFrame(get_number_data_for_df(['A', 'B'], 2), index=pd.RangeIndex(0, 2))
+    python_code = 'df.loc[[0], [\'B\']] = SUM(df[[\'A\']])'
+    functions = set(['SUM'])
+    columns = set(['A'])
+    code, funcs, cols, _ = parse_formula(formula, column_header, formula_label, {'type': FORMULA_SPECIFIC_INDEX_LABELS_TYPE, 'index_labels': [0]}, [df], ['df'], 0) 
+    assert (code, funcs, cols) == \
+        (
+            python_code, 
+            functions, 
+            columns
+        )
+
+def test_specific_index_labels_header_header_multiple_header_dependencies():
+    formula = '=SUM(A:B)'
+    column_header = 'C'
+    formula_label = 0
+    df = pd.DataFrame(get_number_data_for_df(['A', 'B', 'C'], 2), index=pd.RangeIndex(0, 2))
+    python_code = 'df.loc[[0], [\'C\']] = SUM(df.loc[:, \'A\':\'B\'])'
+    functions = set(['SUM'])
+    columns = set(['A', 'B'])
+    code, funcs, cols, _ = parse_formula(formula, column_header, formula_label, {'type': FORMULA_SPECIFIC_INDEX_LABELS_TYPE, 'index_labels': [0]}, [df], ['df'], 0) 
+    assert (code, funcs, cols) == \
+        (
+            python_code, 
+            functions, 
+            columns
+        )
+
 @pytest.mark.parametrize("formula,column_header,formula_label,dfs,df_names,sheet_index,python_code,functions,columns", POST_PD_1_2_VLOOKUP_TESTS)
 @pandas_post_1_2_only
 def post_pandas_1_2_cross_sheet_tests(formula,column_header,formula_label,dfs,df_names,sheet_index,python_code,functions,columns):