dmitri_branching.py

from pendulum import datetime

from airflow.decorators import dag, task, task_group

from airflow.operators.empty import EmptyOperator


@dag(
    tags=["test"],
    schedule=None,
    start_date=datetime(2023, 8, 1),
    catchup=False,
)
def IdillaylaI_branching():
    start = EmptyOperator(task_id='start')
    end = EmptyOperator(task_id='end')

    @task_group
    def gr1():
        pre = EmptyOperator(task_id='prewere')
        calcul1 = EmptyOperator(task_id='calculate1', trigger_rule='one_success')

        @task
        def calculate2():
            return "crunching numbers like there's no tomorrow"

        @task
        def checking():
            return 1

        @task.branch
        def branch(result):
            if not result:
                return 'gr1.prewere'
            else:
                return 'gr1.calculate1'

        branch(checking()) >> [pre, calcul1]
        pre >> calcul1 >> calculate2()

    start >> gr1() >> end


IdillaylaI_branching()