add py op callable fix

Author: akshaykumarsalunke

orbiter_translations/talend/talend_demo.py

@@ -46,7 +46,7 @@ def basic_dag_rule(val: dict) -> OrbiterDAG | None:
         return OrbiterDAG(
             dag_id=dag_id,
             file_path=Path(f"{dag_id}.py"),
-            doc_md=val.get("description", "Created from Talend pipeline")
+            doc_md=val.get("description", "Created from Talend pipeline"),
         )
     return None
 
@@ -55,10 +55,12 @@ def task_common_args(val: dict) -> dict:
     """
     Common mappings for all tasks
     """
-    task_id = (val.get("data", {})
-               .get("properties", {})
-               .get("$componentMetadata", {})
-               .get("name", "UNKNOWN"))
+    task_id = (
+        val.get("data", {})
+        .get("properties", {})
+        .get("$componentMetadata", {})
+        .get("name", "UNKNOWN")
+    )
 
     params = {
         "task_id": task_id.replace(" ", "_"),
@@ -81,27 +83,48 @@ def task_filter_rule(val: dict) -> list[dict] | None:
 @task_rule(priority=2)
 def python_task_rule(val: dict) -> OrbiterPythonOperator | None:
 
-    if "python3" in val.get("data", {}).get("properties", {}).get("$componentMetadata", {}).get("technicalType", "").lower():
-        code = val.get("data", {}).get("properties", {}).get("configuration", {}).get("pythonCode", "")
-        name = val.get("data", {}).get("properties", {}).get("$componentMetadata", {}).get("name", "").lower().replace(" ", "_")
+    if (
+        "python3"
+        in val.get("data", {})
+        .get("properties", {})
+        .get("$componentMetadata", {})
+        .get("technicalType", "")
+        .lower()
+    ):
+        code = (
+            val.get("data", {})
+            .get("properties", {})
+            .get("configuration", {})
+            .get("pythonCode", "")
+        )
+        name = (
+            val.get("data", {})
+            .get("properties", {})
+            .get("$componentMetadata", {})
+            .get("name", "")
+            .lower()
+            .replace(" ", "_")
+        )
         if code and name:
+            indented_code = "\n    ".join([""] + code.splitlines())
+            function_content = f"""def {name}(**context):
+            {indented_code}
+            """
             return OrbiterPythonOperator(
                 orbiter_includes={
-                        OrbiterInclude(
-                            filepath=f"include/{name}.py",
-                            contents=code,
-                        ),
-                    },
-                    imports=[
-                        OrbiterRequirement(
-                            module=f"include.{name}", names=[f"{name}"]
-                        ),
-                        OrbiterRequirement(
-                            module="airflow.operators.python", names=["PythonOperator"]
-                        ),
-                    ],
-                    python_callable=f"{name}",
-                **task_common_args(val)
+                    OrbiterInclude(
+                        filepath=f"include/{name}.py",
+                        contents=function_content,
+                    ),
+                },
+                imports=[
+                    OrbiterRequirement(module=f"include.{name}", names=[f"{name}"]),
+                    OrbiterRequirement(
+                        module="airflow.operators.python", names=["PythonOperator"]
+                    ),
+                ],
+                python_callable=f"{name}",
+                **task_common_args(val),
             )
     return None
 
@@ -124,18 +147,20 @@ def simple_task_dependencies(
 
     component_to_task_name = {}
     for component in pipeline.get("components", []):
-        name = (component.get("data", {})
-                .get("properties", {})
-                .get("$componentMetadata", {})
-                .get("name"))
+        name = (
+            component.get("data", {})
+            .get("properties", {})
+            .get("$componentMetadata", {})
+            .get("name")
+        )
         if name:
             component_to_task_name[component["id"]] = name.replace(" ", "_").lower()
 
     port_to_component = {}
     for port in pipeline.get("ports", []):
-        port_type = (port.get("graphicalAttributes", {})
-                    .get("properties", {})
-                    .get("type"))
+        port_type = (
+            port.get("graphicalAttributes", {}).get("properties", {}).get("type")
+        )
         port_to_component[port["id"]] = (port["nodeId"], port_type)
 
     for step in pipeline.get("steps", []):
@@ -153,10 +178,11 @@ def simple_task_dependencies(
             target_task_name = component_to_task_name.get(target_component_id)
 
             if source_task_name and target_task_name:
-                dependencies.append(OrbiterTaskDependency(
-                    task_id=source_task_name,
-                    downstream=target_task_name
-                ))
+                dependencies.append(
+                    OrbiterTaskDependency(
+                        task_id=source_task_name, downstream=target_task_name
+                    )
+                )
 
     return dependencies if dependencies else None