Fix EnforceTaskFlowApi for task

- do not replace if callable is a lambda
feluelle · Apr 21, 2022 · f86e7b2 · f86e7b2
1 parent 2a79534
commit f86e7b2
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/airflint/rules/task.py b/airflint/rules/task.py
@@ -95,6 +95,10 @@ def match(self, node):
             and isinstance(node.value.func, ast.Name)
             and node.value.func.id in ["PythonOperator", "PythonVirtualenvOperator"]
             and isinstance(node.value.func.ctx, ast.Load)
+            and any(
+                keyword.arg == "python_callable" and isinstance(keyword.value, ast.Name)
+                for keyword in node.value.keywords
+            )
         )
 
         args = next(

diff --git a/tests/test_rules.py b/tests/test_rules.py
@@ -101,6 +101,20 @@ def foo():
             task = PythonOperator(task_id="foo", python_callable=foo)
             """,
         ),
+        (
+            task.EnforceTaskFlowApi,
+            """
+            from airflow.operators.python import PythonOperator
+
+            PythonOperator(task_id="foo", python_callable=lambda: "foo")
+            """,
+            """
+            from airflow.operators.python import PythonOperator
+            from airflow.decorators import task
+
+            PythonOperator(task_id="foo", python_callable=lambda: "foo")
+            """,
+        ),
         (
             task.EnforceTaskFlowApi,
             """