Fix TaskFlowApi and VariableGet replacements

- for task flow api consider only tasks with constants, except for python_callable, op_args and op_kwargs - do not raise an exception on failed import, skip instead
feluelle · Apr 21, 2022 · a447e2a · a447e2a
1 parent af167cd
commit a447e2a
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 1 deletion.
diff --git a/airflint/rules/task.py b/airflint/rules/task.py
@@ -55,6 +55,11 @@ def match(self, node):
                 scope=node_scope,
             )
         )
+        assert all(
+            isinstance(keyword.value, ast.Constant)
+            for keyword in python_operator.keywords
+            if keyword.arg not in ["python_callable", "op_args", "op_kwargs"]
+        )
         assert isinstance(python_operator.func, ast.Name)
         TASK_MAPPING = {
             "PythonOperator": ast.Name(id="task", ctx=ast.Load()),
@@ -100,6 +105,11 @@ def match(self, node):
                 for keyword in node.value.keywords
             )
         )
+        assert all(
+            isinstance(keyword.value, ast.Constant)
+            for keyword in node.value.keywords
+            if keyword.arg not in ["python_callable", "op_args", "op_kwargs"]
+        )
 
         args = next(
             (

diff --git a/airflint/rules/variable.py b/airflint/rules/variable.py
@@ -34,8 +34,12 @@ def match(self, node):
                 )
             )
             and import_node.module
-            and (file_path := importlib.import_module(import_node.module).__file__)
         )
+        try:
+            _module = importlib.import_module(import_node.module)
+        except ImportError:
+            pass
+        assert _module and (file_path := _module.__file__)
         with open(file_path) as file:
             module = ast.parse(file.read())
         assert any(

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