Skip to content

Commit 5b180f2

Browse files
committed
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
1 parent 5c529a9 commit 5b180f2

11 files changed

Lines changed: 82 additions & 167 deletions

File tree

dagfactory/_yaml.py

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
from __future__ import annotations
22

3+
import copy
34
import os
5+
from functools import reduce
46

57
import yaml
68

@@ -12,25 +14,40 @@ def load_yaml_file(file_path: str) -> dict[str, any]:
1214
Load a YAML file into a dictionary.
1315
"""
1416

15-
def __join(loader: yaml.FullLoader, node: yaml.Node) -> str:
16-
seq = loader.construct_sequence(node)
17-
return "".join([str(i) for i in seq])
18-
19-
def __or(loader: yaml.FullLoader, node: yaml.Node) -> str:
20-
seq = loader.construct_sequence(node)
21-
return " | ".join([f"({str(i)})" for i in seq])
22-
23-
def __and(loader: yaml.FullLoader, node: yaml.Node) -> str:
24-
seq = loader.construct_sequence(node)
25-
return " & ".join([f"({str(i)})" for i in seq])
26-
27-
yaml.add_constructor("!join", __join, yaml.FullLoader)
28-
yaml.add_constructor("!or", __or, yaml.FullLoader)
29-
yaml.add_constructor("!and", __and, yaml.FullLoader)
17+
def _flatten_logical_expressions_helper(data):
18+
if isinstance(data, str):
19+
return data
20+
if isinstance(data, dict):
21+
if "__and__" in data:
22+
processed_items = [_flatten_logical_expressions_helper(item) for item in data["__and__"]]
23+
try:
24+
return reduce(lambda a, b: a & b, processed_items)
25+
except TypeError:
26+
return f"({' & '.join(processed_items)})"
27+
if "__or__" in data:
28+
processed_items = [_flatten_logical_expressions_helper(item) for item in data["__or__"]]
29+
try:
30+
return reduce(lambda a, b: a | b, processed_items)
31+
except TypeError:
32+
return f"({' | '.join(processed_items)})"
33+
if "__join__" in data:
34+
processed_items = [_flatten_logical_expressions_helper(item) for item in data["__join__"]]
35+
return "".join(processed_items)
36+
new_dict = {}
37+
for key, value in data.items():
38+
new_dict[key] = _flatten_logical_expressions_helper(value)
39+
return new_dict
40+
if isinstance(data, list):
41+
return [_flatten_logical_expressions_helper(item) for item in data]
42+
return data
43+
44+
def _flatten_logical_expressions(data):
45+
return _flatten_logical_expressions_helper(copy.deepcopy(data))
3046

3147
with open(file_path, "r", encoding="utf-8") as fp:
3248
config_with_env = os.path.expandvars(fp.read())
3349
config: dict[str, any] = yaml.load(stream=config_with_env, Loader=yaml.FullLoader)
3450
config = cast_with_type(config)
51+
config = _flatten_logical_expressions(config)
3552

3653
return config

dagfactory/dagbuilder.py

Lines changed: 7 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@
99
import re
1010
import warnings
1111
from copy import deepcopy
12+
1213
from datetime import datetime
13-
from functools import partial, reduce
14+
from functools import partial
1415
from typing import Any, Callable, Dict, List, Tuple, Union
1516

1617
from airflow import configuration
@@ -592,66 +593,6 @@ def process_file_with_datasets(file: str, datasets_conditions: str) -> Any:
592593
datasets_uri = utils.get_datasets_uri_yaml_file(file, list(dataset_map.keys()))
593594
return [Dataset(uri) for uri in datasets_uri]
594595

595-
@staticmethod
596-
def _init_watchers(watchers_data):
597-
"""Initialize watcher objects from configuration."""
598-
from dagfactory.utils import _import_from_string
599-
600-
watchers = []
601-
for watcher in watchers_data:
602-
watcher_class = _import_from_string(watcher["callable"])
603-
trigger_data = watcher.get("trigger", {})
604-
trigger_class = _import_from_string(trigger_data.get("callable"))
605-
trigger_params = trigger_data.get("params", {})
606-
watchers.append(watcher_class(name=watcher.get("name"), trigger=trigger_class(**trigger_params)))
607-
return watchers
608-
609-
@staticmethod
610-
def _combine_assets(assets, op: str):
611-
"""Combine a list of Asset objects using logical operators."""
612-
if op == "or":
613-
return reduce(lambda a, b: a | b, assets)
614-
elif op == "and":
615-
return reduce(lambda a, b: a & b, assets)
616-
else:
617-
raise ValueError(f"Unknown operator: {op}")
618-
619-
@staticmethod
620-
def _is_asset(d):
621-
from airflow.sdk import Asset
622-
623-
if not isinstance(d, dict):
624-
return False
625-
for key, value in d.items():
626-
if isinstance(value, Asset):
627-
return True
628-
elif isinstance(value, list):
629-
if any(isinstance(item, Asset) for item in value):
630-
return True
631-
elif isinstance(value, dict):
632-
if DagBuilder._is_asset(value):
633-
return True
634-
return False
635-
636-
@staticmethod
637-
def _asset_schedule(value):
638-
"""Recursively parse and construct assets or combinations of assets."""
639-
from airflow.sdk import Asset
640-
641-
if isinstance(value, dict):
642-
if "or" in value:
643-
assets = [DagBuilder._asset_schedule(item) for item in value["or"]]
644-
return DagBuilder._combine_assets(assets, "or")
645-
elif "and" in value:
646-
assets = [DagBuilder._asset_schedule(item) for item in value["and"]]
647-
return DagBuilder._combine_assets(assets, "and")
648-
elif isinstance(value, list):
649-
return [asset for asset in value]
650-
elif isinstance(value, Asset):
651-
return value
652-
else:
653-
raise TypeError(f"Unexpected data type: {type(value)}")
654-
655596
@staticmethod
656597
def configure_schedule(dag_params: Dict[str, Any], dag_kwargs: Dict[str, Any]) -> None:
657598
"""
@@ -722,19 +663,11 @@ def configure_schedule(dag_params: Dict[str, Any], dag_kwargs: Dict[str, Any]) -
722663
if has_datasets_attr:
723664
schedule.pop("datasets")
724665
else:
725-
if "schedule" in dag_params:
726-
schedule = dag_params.get("schedule")
727-
if DagBuilder._is_asset(schedule):
728-
dag_kwargs[schedule_key] = DagBuilder._asset_schedule(schedule)
729-
else:
730-
if (
731-
utils.check_dict_key(dag_params, "schedule")
732-
and isinstance(dag_params["schedule"], str)
733-
and dag_params["schedule"].strip().lower() == "none"
734-
):
735-
dag_kwargs[schedule_key] = None
736-
else:
737-
dag_kwargs[schedule_key] = schedule
666+
schedule = dag_params.get("schedule")
667+
if utils.check_dict_key(dag_params, "schedule") and isinstance(schedule, str) and schedule.strip().lower() == "none":
668+
dag_kwargs[schedule_key] = None
669+
else:
670+
dag_kwargs[schedule_key] = dag_params.get("schedule")
738671

739672
@staticmethod
740673
def _normalise_tasks_config(tasks_cfg: Any) -> Dict[str, Dict[str, Any]]:

dagfactory/dagfactory.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,11 @@ def __init__(
5555
DagFactory._validate_config_filepath(config_filepath=config_filepath)
5656
self.config: Dict[str, Any] = self._load_dag_config(config_filepath=config_filepath)
5757
if config:
58-
self.config: Dict[str, Any] = config
58+
self.config = config
59+
# This will only invoke in the CI
60+
# Make yaml DAG compatible for Airflow 3
61+
if version.parse(AIRFLOW_VERSION) >= version.parse("3.0.0") and os.getenv("AUTO_CONVERT_TO_AF3"):
62+
self.config = update_yaml_structure(config)
5963

6064
# These default args are a bit different; these are not the "default" structure that is applied to certain DAGs.
6165
# These are in-fact the "default" default_args

dev/dags/datasets/example_dag_datasets.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,8 @@ example_without_custom_config_condition_dataset_consumer_dag:
6868
catchup: false
6969
schedule:
7070
datasets:
71-
!or
72-
- !and
71+
__or__:
72+
- __and__:
7373
- "s3://bucket-cjmm/raw/dataset_custom_1"
7474
- "s3://bucket-cjmm/raw/dataset_custom_2"
7575
- "s3://bucket-cjmm/raw/dataset_custom_3"

dev/dags/datasets/example_dataset_yaml_syntax.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ consumer_dag:
66
catchup: false
77
schedule:
88
datasets:
9-
!or
10-
- !and
9+
__or__:
10+
- __and__:
1111
- "s3://bucket-cjmm/raw/dataset_custom_1"
1212
- "s3://bucket-cjmm/raw/dataset_custom_2"
1313
- "s3://bucket-cjmm/raw/dataset_custom_3"

tests/fixtures/dag_factory.yml

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,13 @@ example_dag3:
6666
example_dag4:
6767
vars:
6868
arg1: &arg1 'hello'
69-
arg2: &arg2 !join [*arg1, ' world']
70-
tasks:
71-
- task_id: "task_1"
72-
bash_command: !join ['echo ', *arg2]
73-
operator: airflow.operators.bash.BashOperator
69+
arg2: &arg2
70+
__join__:
71+
- *arg1
72+
- ' world'
73+
tasks:
74+
- task_id: task_1
75+
bash_command:
76+
__join__:
77+
- 'echo '
78+
- *arg2

tests/fixtures_without_default_yaml/dag_factory.yml

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,14 @@ example_dag3:
6060
example_dag4:
6161
vars:
6262
arg1: &arg1 'hello'
63-
arg2: &arg2 !join [*arg1, ' world']
63+
arg2: &arg2
64+
__join__:
65+
- *arg1
66+
- ' world'
6467
tasks:
6568
- task_id: "task_1"
66-
bash_command: !join ['echo ', *arg2]
69+
bash_command:
70+
__join__:
71+
- 'echo '
72+
- *arg2
6773
operator: airflow.operators.bash.BashOperator

tests/schedule/and_asset.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
schedule:
2-
and:
2+
__and__:
33
- __type__: airflow.sdk.Asset
44
uri: s3://dag1/output_1.txt
55
extra:

tests/schedule/nested_asset.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
schedule:
2-
or:
3-
- and:
2+
__or__:
3+
- __and__:
44
- __type__: airflow.sdk.Asset
55
uri: s3://dag1/output_1.txt
66
extra:

tests/schedule/or_asset.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
schedule:
2-
or:
2+
__or__:
33
- __type__: airflow.sdk.Asset
44
uri: s3://dag1/output_1.txt
55
extra:

0 commit comments

Comments
 (0)