Remove custom yml constructor and consolidate boolean implementation

Following up on PR #487, this PR demonstrates and replaces the usages of
tasks and task groups across all example dags, tests and static yamls in
docs to be a list.

closes: #491

Fix yml

Author: pankajastro

dagfactory/_yaml.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import copy
 import os
+from functools import reduce
 
 import yaml
 
@@ -12,25 +14,40 @@ def load_yaml_file(file_path: str) -> dict[str, any]:
     Load a YAML file into a dictionary.
     """
 
-    def __join(loader: yaml.FullLoader, node: yaml.Node) -> str:
-        seq = loader.construct_sequence(node)
-        return "".join([str(i) for i in seq])
-
-    def __or(loader: yaml.FullLoader, node: yaml.Node) -> str:
-        seq = loader.construct_sequence(node)
-        return " | ".join([f"({str(i)})" for i in seq])
-
-    def __and(loader: yaml.FullLoader, node: yaml.Node) -> str:
-        seq = loader.construct_sequence(node)
-        return " & ".join([f"({str(i)})" for i in seq])
-
-    yaml.add_constructor("!join", __join, yaml.FullLoader)
-    yaml.add_constructor("!or", __or, yaml.FullLoader)
-    yaml.add_constructor("!and", __and, yaml.FullLoader)
+    def _flatten_logical_expressions_helper(data):
+        if isinstance(data, str):
+            return data
+        if isinstance(data, dict):
+            if "__and__" in data:
+                processed_items = [_flatten_logical_expressions_helper(item) for item in data["__and__"]]
+                try:
+                    return reduce(lambda a, b: a & b, processed_items)
+                except TypeError:
+                    return f"({' & '.join(processed_items)})"
+            if "__or__" in data:
+                processed_items = [_flatten_logical_expressions_helper(item) for item in data["__or__"]]
+                try:
+                    return reduce(lambda a, b: a | b, processed_items)
+                except TypeError:
+                    return f"({' | '.join(processed_items)})"
+            if "__join__" in data:
+                processed_items = [_flatten_logical_expressions_helper(item) for item in data["__join__"]]
+                return "".join(processed_items)
+            new_dict = {}
+            for key, value in data.items():
+                new_dict[key] = _flatten_logical_expressions_helper(value)
+            return new_dict
+        if isinstance(data, list):
+            return [_flatten_logical_expressions_helper(item) for item in data]
+        return data
+
+    def _flatten_logical_expressions(data):
+        return _flatten_logical_expressions_helper(copy.deepcopy(data))
 
     with open(file_path, "r", encoding="utf-8") as fp:
         config_with_env = os.path.expandvars(fp.read())
         config: dict[str, any] = yaml.load(stream=config_with_env, Loader=yaml.FullLoader)
         config = cast_with_type(config)
+        config = _flatten_logical_expressions(config)
 
     return config

dagfactory/dagbuilder.py

@@ -9,8 +9,9 @@
 import re
 import warnings
 from copy import deepcopy
+
 from datetime import datetime
-from functools import partial, reduce
+from functools import partial
 from typing import Any, Callable, Dict, List, Tuple, Union
 
 from airflow import configuration
@@ -594,66 +595,6 @@ def process_file_with_datasets(file: str, datasets_conditions: str) -> Any:
             datasets_uri = utils.get_datasets_uri_yaml_file(file, list(dataset_map.keys()))
             return [Dataset(uri) for uri in datasets_uri]
 
-    @staticmethod
-    def _init_watchers(watchers_data):
-        """Initialize watcher objects from configuration."""
-        from dagfactory.utils import _import_from_string
-
-        watchers = []
-        for watcher in watchers_data:
-            watcher_class = _import_from_string(watcher["callable"])
-            trigger_data = watcher.get("trigger", {})
-            trigger_class = _import_from_string(trigger_data.get("callable"))
-            trigger_params = trigger_data.get("params", {})
-            watchers.append(watcher_class(name=watcher.get("name"), trigger=trigger_class(**trigger_params)))
-        return watchers
-
-    @staticmethod
-    def _combine_assets(assets, op: str):
-        """Combine a list of Asset objects using logical operators."""
-        if op == "or":
-            return reduce(lambda a, b: a | b, assets)
-        elif op == "and":
-            return reduce(lambda a, b: a & b, assets)
-        else:
-            raise ValueError(f"Unknown operator: {op}")
-
-    @staticmethod
-    def _is_asset(d):
-        from airflow.sdk import Asset
-
-        if not isinstance(d, dict):
-            return False
-        for key, value in d.items():
-            if isinstance(value, Asset):
-                return True
-            elif isinstance(value, list):
-                if any(isinstance(item, Asset) for item in value):
-                    return True
-            elif isinstance(value, dict):
-                if DagBuilder._is_asset(value):
-                    return True
-        return False
-
-    @staticmethod
-    def _asset_schedule(value):
-        """Recursively parse and construct assets or combinations of assets."""
-        from airflow.sdk import Asset
-
-        if isinstance(value, dict):
-            if "or" in value:
-                assets = [DagBuilder._asset_schedule(item) for item in value["or"]]
-                return DagBuilder._combine_assets(assets, "or")
-            elif "and" in value:
-                assets = [DagBuilder._asset_schedule(item) for item in value["and"]]
-                return DagBuilder._combine_assets(assets, "and")
-        elif isinstance(value, list):
-            return [asset for asset in value]
-        elif isinstance(value, Asset):
-            return value
-        else:
-            raise TypeError(f"Unexpected data type: {type(value)}")
-
     @staticmethod
     def configure_schedule(dag_params: Dict[str, Any], dag_kwargs: Dict[str, Any]) -> None:
         """
@@ -700,17 +641,14 @@ def configure_schedule(dag_params: Dict[str, Any], dag_kwargs: Dict[str, Any]) -
                     schedule.pop("datasets")
         else:
             schedule = dag_params.get("schedule")
-            if DagBuilder._is_asset(schedule):
-                dag_kwargs["schedule"] = DagBuilder._asset_schedule(schedule)
+            if (
+                utils.check_dict_key(dag_params, "schedule")
+                and isinstance(dag_params["schedule"], str)
+                and dag_params["schedule"].strip().lower() == "none"
+            ):
+                dag_kwargs["schedule"] = None
             else:
-                if (
-                    utils.check_dict_key(dag_params, "schedule")
-                    and isinstance(dag_params["schedule"], str)
-                    and dag_params["schedule"].strip().lower() == "none"
-                ):
-                    dag_kwargs["schedule"] = None
-                else:
-                    dag_kwargs["schedule"] = schedule
+                dag_kwargs["schedule"] = schedule
 
     @staticmethod
     def _normalise_tasks_config(tasks_cfg: Any) -> Dict[str, Dict[str, Any]]:

dagfactory/dagfactory.py

@@ -55,7 +55,11 @@ def __init__(
             DagFactory._validate_config_filepath(config_filepath=config_filepath)
             self.config: Dict[str, Any] = self._load_dag_config(config_filepath=config_filepath)
         if config:
-            self.config: Dict[str, Any] = config
+            self.config = config
+            # This will only invoke in the CI
+            # Make yaml DAG compatible for Airflow 3
+            if version.parse(AIRFLOW_VERSION) >= version.parse("3.0.0") and os.getenv("AUTO_CONVERT_TO_AF3"):
+                self.config = update_yaml_structure(config)
 
         # These default args are a bit different; these are not the "default" structure that is applied to certain DAGs.
         # These are in-fact the "default" default_args

dev/dags/datasets/example_dag_datasets.yml

@@ -68,8 +68,8 @@ example_without_custom_config_condition_dataset_consumer_dag:
   catchup: false
   schedule:
     datasets:
-      !or
-        - !and
+      __or__:
+        - __and__:
           - "s3://bucket-cjmm/raw/dataset_custom_1"
           - "s3://bucket-cjmm/raw/dataset_custom_2"
         - "s3://bucket-cjmm/raw/dataset_custom_3"

dev/dags/datasets/example_dataset_yaml_syntax.yml

@@ -6,8 +6,8 @@ consumer_dag:
   catchup: false
   schedule:
     datasets:
-      !or
-        - !and
+      __or__:
+        - __and__:
           - "s3://bucket-cjmm/raw/dataset_custom_1"
           - "s3://bucket-cjmm/raw/dataset_custom_2"
         - "s3://bucket-cjmm/raw/dataset_custom_3"

tests/fixtures/dag_factory.yml

@@ -66,8 +66,13 @@ example_dag3:
 example_dag4:
   vars:
     arg1: &arg1 'hello'
-    arg2: &arg2 !join [*arg1, ' world']
-  tasks:
-    - task_id: "task_1"
-      bash_command: !join ['echo ', *arg2]
-      operator: airflow.operators.bash.BashOperator
+    arg2: &arg2
+      __join__:
+        - *arg1
+        - ' world'
+    tasks:
+      - task_id: task_1
+        bash_command:
+          __join__:
+            - 'echo '
+            - *arg2

tests/fixtures_without_default_yaml/dag_factory.yml

@@ -60,8 +60,14 @@ example_dag3:
 example_dag4:
   vars:
     arg1: &arg1 'hello'
-    arg2: &arg2 !join [*arg1, ' world']
+    arg2: &arg2
+      __join__:
+        - *arg1
+        - ' world'
   tasks:
     - task_id: "task_1"
-      bash_command: !join ['echo ', *arg2]
+      bash_command:
+        __join__:
+          - 'echo '
+          - *arg2
       operator: airflow.operators.bash.BashOperator

tests/schedule/and_asset.yml

@@ -1,5 +1,5 @@
 schedule:
-  and:
+  __and__:
     - __type__: airflow.sdk.Asset
       uri: s3://dag1/output_1.txt
       extra:

tests/schedule/nested_asset.yml

@@ -1,6 +1,6 @@
 schedule:
-  or:
-    - and:
+  __or__:
+    - __and__:
         - __type__: airflow.sdk.Asset
           uri: s3://dag1/output_1.txt
           extra:

tests/schedule/or_asset.yml

@@ -1,5 +1,5 @@
 schedule:
-  or:
+  __or__:
     - __type__: airflow.sdk.Asset
       uri: s3://dag1/output_1.txt
       extra: