[sk] Bug fixes for transformer action template generation (mage-ai#653)

Aliang-CN · Jul 14, 2022 · ea14c75 · ea14c75
1 parent 5fcd8e8
commit ea14c75
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/mage_ai/data_cleaner/transformer_actions/utils.py b/mage_ai/data_cleaner/transformer_actions/utils.py
@@ -75,7 +75,7 @@ def build_transformer_action(
     Returns:
         Dict: Transformer action payload
     """
-    action_variables = build_action_variables(df, action_arguments)
+    action_variables = build_action_variables(df)
     return dict(
         action_type=action_type,
         action_arguments=action_arguments,

diff --git a/mage_ai/data_preparation/templates/transformers/transformer_action_fmt.jinja b/mage_ai/data_preparation/templates/transformers/transformer_action_fmt.jinja
@@ -13,6 +13,7 @@ def execute_transformer_action(df: DataFrame, *args) -> DataFrame:
     Execute Transformer Action: {{ action_type }}
     """
     action = build_transformer_action(
+        df,
         action_type=ActionType.{{ action_type.upper() }},
         action_arguments=[],
         axis=Axis.{{ axis.upper() }},{{ kwargs }}

diff --git a/mage_ai/tests/data_preparation/test_templates.py b/mage_ai/tests/data_preparation/test_templates.py
@@ -242,6 +242,7 @@ def execute_transformer_action(df: DataFrame, *args) -> DataFrame:
     Execute Transformer Action: clean_column_name
     \"\"\"
     action = build_transformer_action(
+        df,
         action_type=ActionType.CLEAN_COLUMN_NAME,
         action_arguments=[],
         axis=Axis.COLUMN,
@@ -270,6 +271,7 @@ def execute_transformer_action(df: DataFrame, *args) -> DataFrame:
     Execute Transformer Action: custom
     \"\"\"
     action = build_transformer_action(
+        df,
         action_type=ActionType.CUSTOM,
         action_arguments=[],
         axis=Axis.ROW,
@@ -299,6 +301,7 @@ def execute_transformer_action(df: DataFrame, *args) -> DataFrame:
     Execute Transformer Action: reformat
     \"\"\"
     action = build_transformer_action(
+        df,
         action_type=ActionType.REFORMAT,
         action_arguments=[],
         axis=Axis.COLUMN,
@@ -328,6 +331,7 @@ def execute_transformer_action(df: DataFrame, *args) -> DataFrame:
     Execute Transformer Action: add
     \"\"\"
     action = build_transformer_action(
+        df,
         action_type=ActionType.ADD,
         action_arguments=[],
         axis=Axis.COLUMN,