multimeric opened a new issue, #61758: URL: https://github.com/apache/airflow/issues/61758
### Description Task groups are allowed to return a `DAGNode`: https://github.com/apache/airflow/blob/752419dae9d27419019395b1aaf9ea1bf497b5e9/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py#L204-L206 https://github.com/apache/airflow/blob/752419dae9d27419019395b1aaf9ea1bf497b5e9/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py#L55 A typical `DAGNode` is `Operator`. However in taskflow you typically work with `XComArg`, not `Operator`. ```python from typing import reveal_type from airflow.sdk import dag, task, task_group @task_group def some_group() -> Operator: @task def generate_data() -> list[int]: return [1, 2, 3, 4, 5] result = generate_data() reveal_type(result) ``` ``` XComArg ``` To pull a `DAGNode` out from taskflow, you could do something like this, but it's digging into the internals: ```python if isinstance(result, PlainXComArg): return result.operator else: raise ValueError("Expected a PlainXComArg, got something else.") ``` The best solution would be to allow the `@task_group` decorated function to return an `XComArg`, and then the `DAGNode` could be inferred from there. Failing that, just a consistent API for retrieving the operator from an `XComArg` would be ideal. *** The other issue is at the other end of the task group: when you use it. ```python @dag( start_date=datetime(2024, 1, 1), ) def workflow(): group_result = some_group() reveal_type(group_result) ``` ``` DAGNode ``` This is the opposite problem. Here we have a `DAGNode`, but taskflow wants to deal with `XComArg`. The hacky workaround is to cast it to `Operator` and then use `.output`, but this is the same problem. So I think the `@task_group` decorated function would ideally return `XComArg` as well. ### Use case/motivation This would effectively allow you to use the return value of a `@task_group`. Currently task groups work as a kind of subDAG that can't return values. ### Related issues _No response_ ### Are you willing to submit a PR? - [x] Yes I am willing to submit a PR! ### Code of Conduct - [x] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
