Add support for python_callable_file for BranchPythonOperator (astronomer#84)

* Add support for python_callable_file for BranchPythonOperator

* Fix pylint errors in dagbuilder

* Reformat using black

Co-authored-by: Aditya Aggarwal <[email protected]>

Author: the-cybersapien

dagfactory/dagbuilder.py

@@ -9,7 +9,7 @@
 from airflow.models import Variable
 from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator
 from airflow.models import BaseOperator
-from airflow.operators.python_operator import PythonOperator
+from airflow.operators.python_operator import PythonOperator, BranchPythonOperator
 from airflow.utils.module_loading import import_string
 from airflow import __version__ as AIRFLOW_VERSION
 
@@ -179,13 +179,13 @@ def make_task(operator: str, task_params: Dict[str, Any]) -> BaseOperator:
         except Exception as err:
             raise Exception(f"Failed to import operator: {operator}") from err
         try:
-            if operator_obj == PythonOperator:
+            if operator_obj in [PythonOperator, BranchPythonOperator]:
                 if not task_params.get("python_callable_name") and not task_params.get(
                     "python_callable_file"
                 ):
                     raise Exception(
-                        "Failed to create task. PythonOperator requires `python_callable_name` \
-                        and `python_callable_file` parameters."
+                        "Failed to create task. PythonOperator and BranchPythonOperator requires \
+                        `python_callable_name` and `python_callable_file` parameters."
                     )
                 task_params["python_callable"]: Callable = utils.get_python_callable(
                     task_params["python_callable_name"],

examples/customized/operators/breakfast_operators.py

@@ -2,7 +2,7 @@
 
 
 class MakeBreadOperator(BaseOperator):
-    template_fields = ('bread_type',)
+    template_fields = ("bread_type",)
 
     def __init__(self, bread_type, *args, **kwargs):
         super(MakeBreadOperator, self).__init__(*args, **kwargs)
@@ -13,7 +13,7 @@ def execute(self, context):
 
 
 class MakeCoffeeOperator(BaseOperator):
-    template_fields = ('coffee_type',)
+    template_fields = ("coffee_type",)
 
     def __init__(self, coffee_type, *args, **kwargs):
         super(MakeCoffeeOperator, self).__init__(*args, **kwargs)

tests/fixtures/doc_md_builder.py

@@ -1,2 +1,2 @@
 def mydocmdbuilder(**kwargs):
-	return f"{kwargs}"
+    return f"{kwargs}"

tests/test_dagbuilder.py

@@ -33,12 +33,12 @@
     "default_args": {"owner": "custom_owner"},
     "description": "this is an example dag",
     "schedule_interval": "0 3 * * *",
-    "tags" : ["tag1","tag2"],
+    "tags": ["tag1", "tag2"],
     "tasks": {
         "task_1": {
             "operator": "airflow.operators.bash_operator.BashOperator",
             "bash_command": "echo 1",
-            "execution_timeout_secs" : 5,
+            "execution_timeout_secs": 5,
         },
         "task_2": {
             "operator": "airflow.operators.bash_operator.BashOperator",
@@ -63,7 +63,7 @@
         "task_group_2": {
             "dependencies": ["task_group_1"],
         },
-        "task_group_3": {}
+        "task_group_3": {},
     },
     "tasks": {
         "task_1": {
@@ -110,15 +110,15 @@
     },
     "description": "this is an example dag",
     "schedule_interval": "0 3 * * *",
-    "tags" : ["tag1","tag2"],
+    "tags": ["tag1", "tag2"],
     "on_failure_callback": f"{__name__}.print_context_callback",
     "on_success_callback": f"{__name__}.print_context_callback",
     "sla_miss_callback": f"{__name__}.print_context_callback",
     "tasks": {
         "task_1": {
             "operator": "airflow.operators.bash_operator.BashOperator",
             "bash_command": "echo 1",
-            "execution_timeout_secs" : 5,
+            "execution_timeout_secs": 5,
             "on_failure_callback": f"{__name__}.print_context_callback",
             "on_success_callback": f"{__name__}.print_context_callback",
             "on_execute_callback": f"{__name__}.print_context_callback",
@@ -175,7 +175,7 @@ def test_get_dag_params():
             "task_1": {
                 "operator": "airflow.operators.bash_operator.BashOperator",
                 "bash_command": "echo 1",
-                "execution_timeout_secs": 5
+                "execution_timeout_secs": 5,
             },
             "task_2": {
                 "operator": "airflow.operators.bash_operator.BashOperator",
@@ -202,7 +202,11 @@ def test_get_dag_params_no_start_date():
 def test_make_task_valid():
     td = dagbuilder.DagBuilder("test_dag", DAG_CONFIG, DEFAULT_CONFIG)
     operator = "airflow.operators.bash_operator.BashOperator"
-    task_params = {"task_id": "test_task", "bash_command": "echo 1","execution_timeout_secs":5}
+    task_params = {
+        "task_id": "test_task",
+        "bash_command": "echo 1",
+        "execution_timeout_secs": 5,
+    }
     actual = td.make_task(operator, task_params)
     assert actual.task_id == "test_task"
     assert actual.bash_command == "echo 1"
@@ -258,8 +262,8 @@ def test_build():
     assert isinstance(actual["dag"], DAG)
     assert len(actual["dag"].tasks) == 3
     assert actual["dag"].task_dict["task_1"].downstream_task_ids == {"task_2", "task_3"}
-    if version.parse(AIRFLOW_VERSION) >= version.parse('1.10.8') :
-        assert actual["dag"].tags == ["tag1","tag2"]
+    if version.parse(AIRFLOW_VERSION) >= version.parse("1.10.8"):
+        assert actual["dag"].tags == ["tag1", "tag2"]
 
 
 def test_get_dag_params():
@@ -274,7 +278,10 @@ def test_get_dag_params():
         },
         "schedule_interval": "0 3 * * *",
         "task_groups": {
-            "task_group_1": {"tooltip": "this is a task group", "dependencies": ["task_1"]},
+            "task_group_1": {
+                "tooltip": "this is a task group",
+                "dependencies": ["task_1"],
+            },
             "task_group_2": {"dependencies": ["task_group_1"]},
             "task_group_3": {},
         },
@@ -332,17 +339,24 @@ def test_build_task_groups():
             td.build()
     else:
         actual = td.build()
-        task_group_1 = {t for t in actual["dag"].task_dict if t.startswith("task_group_1")}
-        task_group_2 = {t for t in actual["dag"].task_dict if t.startswith("task_group_2")}
+        task_group_1 = {
+            t for t in actual["dag"].task_dict if t.startswith("task_group_1")
+        }
+        task_group_2 = {
+            t for t in actual["dag"].task_dict if t.startswith("task_group_2")
+        }
         assert actual["dag_id"] == "test_dag"
         assert isinstance(actual["dag"], DAG)
         assert len(actual["dag"].tasks) == 6
-        assert actual["dag"].task_dict["task_1"].downstream_task_ids == {"task_group_1.task_2"}
+        assert actual["dag"].task_dict["task_1"].downstream_task_ids == {
+            "task_group_1.task_2"
+        }
         assert actual["dag"].task_dict["task_group_1.task_2"].downstream_task_ids == {
             "task_group_1.task_3"
         }
         assert actual["dag"].task_dict["task_group_1.task_3"].downstream_task_ids == {
-            "task_4", "task_group_2.task_5",
+            "task_4",
+            "task_group_2.task_5",
         }
         assert actual["dag"].task_dict["task_group_2.task_5"].downstream_task_ids == {
             "task_group_2.task_6",
@@ -360,7 +374,9 @@ def test_make_task_groups():
     }
     dag = "dag"
     task_groups = dagbuilder.DagBuilder.make_task_groups(task_group_dict, dag)
-    expected = MockTaskGroup(tooltip="this is a task group", group_id="task_group", dag=dag)
+    expected = MockTaskGroup(
+        tooltip="this is a task group", group_id="task_group", dag=dag
+    )
     if version.parse(AIRFLOW_VERSION) < version.parse("2.0.0"):
         assert task_groups == {}
     else:
@@ -371,6 +387,7 @@ def test_make_task_groups_empty():
     task_groups = dagbuilder.DagBuilder.make_task_groups({}, None)
     assert task_groups == {}
 
+
 def print_context_callback(context, **kwargs):
     print(context)
 
@@ -393,10 +410,11 @@ def test_make_task_with_callback():
     assert isinstance(actual, PythonOperator)
     assert callable(actual.on_failure_callback)
     assert callable(actual.on_success_callback)
-    if version.parse(AIRFLOW_VERSION) >= version.parse('2.0.0') :
+    if version.parse(AIRFLOW_VERSION) >= version.parse("2.0.0"):
         assert callable(actual.on_execute_callback)
     assert callable(actual.on_retry_callback)
 
+
 def test_make_dag_with_callback():
     td = dagbuilder.DagBuilder("test_dag", DAG_CONFIG_CALLBACK, DEFAULT_CONFIG)
-    td.build()
+    td.build()