def set_work_task_roots_and_leaves()

in airflow-core/src/airflow/utils/setup_teardown.py [0:0]

• 35 lines of code
• 25 McCabe index (conditional complexity)

    def set_work_task_roots_and_leaves(cls):
        """Set the work task roots and leaves."""
        if setup_task := cls._context_managed_setup_task:
            if isinstance(setup_task, list):
                setup_task = tuple(setup_task)
            tasks_in_context = [
                x for x in cls.context_map.get(setup_task, []) if not x.is_teardown and not x.is_setup
            ]
            if tasks_in_context:
                roots = [task for task in tasks_in_context if not task.upstream_list]
                if not roots:
                    setup_task >> tasks_in_context[0]
                else:
                    cls.set_dependency(roots, setup_task, upstream=False)
                leaves = [task for task in tasks_in_context if not task.downstream_list]
                if not leaves:
                    leaves = tasks_in_context[-1]
                cls.set_teardown_task_as_leaves(leaves)

        if teardown_task := cls._context_managed_teardown_task:
            if isinstance(teardown_task, list):
                teardown_task = tuple(teardown_task)
            tasks_in_context = [
                x for x in cls.context_map.get(teardown_task, []) if not x.is_teardown and not x.is_setup
            ]
            if tasks_in_context:
                leaves = [task for task in tasks_in_context if not task.downstream_list]
                if not leaves:
                    teardown_task << tasks_in_context[-1]
                else:
                    cls.set_dependency(leaves, teardown_task)
                roots = [task for task in tasks_in_context if not task.upstream_list]
                if not roots:
                    roots = tasks_in_context[0]
                cls.set_setup_task_as_roots(roots)
        cls.set_setup_teardown_relationships()
        cls.active = False