@task_group returns int, but it appears in @task as TaskGroup

Apache Airflow version

13faa6912f7cd927737a1dc15630d3bbaf2f5d4d

Environment

• Configuration: Local Executor
• OS (e.g. from /etc/os-release): Mac OS 11.3
• Kernel: Darwin Kernel Version 20.4.0
• Install tools: pip install -e .

The DAG

@task
def one():
    return 1

@task_group
def trivial_group(inval):

    @task
    def add_one(i):
        return i + 1

    outval = add_one(inval)
    return outval

@task
def print_it(inval):
    print(inval)

@dag(schedule_interval=None, start_date=days_ago(1), default_args={"owner": "airflow"})
def wrap():

    x = one()
    y = trivial_group(x)
    z = print_it(y)

wrap_dag = wrap()

What happened: print_it had no predecessors and receives <airflow.utils.task_group.TaskGroup object at 0x128921940>

What you expected to happen: print_it comes after trivial_group.add_one and receives 2

The caller ends up with the task group itself, equivalent in the traditional api to tg_ref in:

with TaskGroup("trivial_group") tg_ref:
    pass

This interrupts the ability to continue using the Task Flow api because passing it into a function annotated with @task fails to register the dependency with whatever magic gets it out of xcom and adds edges to the dag.

To Replicate

 $ airflow dags test wrap $(date "+%Y-%m-%d")

[ashb]

Yeah, not released yet. (Sorry, I checked before adding the high priority label but forgot to comment)

We’ll change this before 2.1 is cut

[MatrixManAtYrService]

@ashb the @task_group decorator was merged here: https://github.com/apache/airflow/pull/15034 (Apr 1).

I think that means that it’s not in any released packages. DAG’s using it would have opted in to the bleeding edge by cloning master instead of waiting for a release. Seems reasonable to fix it directly.