Repository: apache/airflow

Issue:
Error setting dependencies on task_group defined using the decorator
### Apache Airflow version

2.2.2 (latest released)

### Operating System

MacOS 11.6.1

### Versions of Apache Airflow Providers

$ pip freeze | grep airflow
apache-airflow==2.2.2
apache-airflow-providers-celery==2.1.0
apache-airflow-providers-ftp==2.0.1
apache-airflow-providers-http==2.0.1
apache-airflow-providers-imap==2.0.1
apache-airflow-providers-sqlite==2.0.1

### Deployment

Other

### Deployment details

`airflow standalone`

### What happened

```
AttributeError: 'NoneType' object has no attribute 'update_relative'
```

### What you expected to happen

Task group should be set as downstream of `start` task, and upstream of `end` task

### How to reproduce

* Add the following code to dags folder

```python
from datetime import datetime

from airflow.decorators import dag, task, task_group


@dag(start_date=datetime(2023, 1, 1), schedule_interval="@once")
def test_dag_1():
    @task
    def start():
        pass

    @task
    def do_thing(x):
        print(x)

    @task_group
    def do_all_things():
        do_thing(1)
        do_thing(2)

    @task
    def end():
        pass

    start() >> do_all_things() >> end()


test_dag_1()
```

* Run `airflow standalone`

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] 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)



Discussion:
Looking at this issue. 
The problem is here, for some when the right shift is called for start, the other value is None, its almost like the task group is not created properly.

```
    def __rshift__(self, other: Union["DependencyMixin", Sequence["DependencyMixin"]]):
        """Implements Task >> Task"""
        self.set_downstream(other)
        return other
```

```
        if not isinstance(task_or_task_list, Sequence):
            task_or_task_list = [task_or_task_list]

        task_list: List["BaseOperator"] = []
        for task_object in task_or_task_list:
```
Looks like the problem is here for both the tasks, with pass or print, the task instances are not created, but replacing it with sleep, a task instance is created and added to the DB. 
` @task
    def end():
        pass`
`@task
    def end():
        pass`

It might be useful to indicate this to the user, still trying to figure out where this logic is performed.

I don’t recall any logic checking a function’s content, so this is very surprising. I can’t reproduce though; what’s a minimal example of this happening?

This seems to work correctly for me:

```python
import datetime

from airflow.decorators import dag, task


@dag(start_date=datetime.datetime(2022, 1, 1))
def empty_task_dag():
    @task
    def empty_task():
        pass

    empty_task()


dag = empty_task_dag()
```

A task instance of DAG ID `empty_task_dag` and task ID `empty_task` is shown on the UI when I manually trigger a DAG run.
Im able to reproduce this in `standalone` mode even without the taskgroup ` start() >> end()`

```
@dag(start_date=datetime(2023, 1, 1), schedule_interval="@once")
def test_dag_1():
    @task
    def start():
        pass

    @task
    def do_thing(x):
        print(x)

    @task_group
    def do_all_things():
        do_thing(1)
        do_thing(2)

    @task
    def end():
        pass

    start() >> do_all_things() >> end()


test_dag_1()
```


Resolving patch:
diff --git a/airflow/decorators/task_group.py b/airflow/decorators/task_group.py
--- a/airflow/decorators/task_group.py
+++ b/airflow/decorators/task_group.py
@@ -42,7 +42,7 @@
 class TaskGroupDecorator(Generic[R]):
     """:meta private:"""
 
-    function: Callable[..., R] = attr.ib(validator=attr.validators.is_callable())
+    function: Callable[..., Optional[R]] = attr.ib(validator=attr.validators.is_callable())
     kwargs: Dict[str, Any] = attr.ib(factory=dict)
     """kwargs for the TaskGroup"""
 
@@ -62,9 +62,24 @@ def _make_task_group(self, **kwargs) -> TaskGroup:
         return TaskGroup(**kwargs)
 
     def __call__(self, *args, **kwargs) -> R:
-        with self._make_task_group(add_suffix_on_collision=True, **self.kwargs):
+        with self._make_task_group(add_suffix_on_collision=True, **self.kwargs) as task_group:
             # Invoke function to run Tasks inside the TaskGroup
-            return self.function(*args, **kwargs)
+            retval = self.function(*args, **kwargs)
+
+        # If the task-creating function returns a task, forward the return value
+        # so dependencies bind to it. This is equivalent to
+        #   with TaskGroup(...) as tg:
+        #       t2 = task_2(task_1())
+        #   start >> t2 >> end
+        if retval is not None:
+            return retval
+
+        # Otherwise return the task group as a whole, equivalent to
+        #   with TaskGroup(...) as tg:
+        #       task_1()
+        #       task_2()
+        #   start >> tg >> end
+        return task_group
 
     def partial(self, **kwargs) -> "MappedTaskGroupDecorator[R]":
         return MappedTaskGroupDecorator(function=self.function, kwargs=self.kwargs).partial(**kwargs)


Test patch:
diff --git a/tests/utils/test_task_group.py b/tests/utils/test_task_group.py
--- a/tests/utils/test_task_group.py
+++ b/tests/utils/test_task_group.py
@@ -689,6 +689,42 @@ def section_2(value2):
     assert extract_node_id(task_group_to_dict(dag.task_group)) == node_ids
 
 
+def test_build_task_group_depended_by_task():
+    """A decorator-based task group should be able to be used as a relative to operators."""
+
+    from airflow.decorators import dag as dag_decorator, task
+
+    @dag_decorator(start_date=pendulum.now())
+    def build_task_group_depended_by_task():
+        @task
+        def task_start():
+            return "[Task_start]"
+
+        @task
+        def task_end():
+            return "[Task_end]"
+
+        @task
+        def task_thing(value):
+            return f"[Task_thing {value}]"
+
+        @task_group_decorator
+        def section_1():
+            task_thing(1)
+            task_thing(2)
+
+        task_start() >> section_1() >> task_end()
+
+    dag = build_task_group_depended_by_task()
+    task_thing_1 = dag.task_dict["section_1.task_thing"]
+    task_thing_2 = dag.task_dict["section_1.task_thing__1"]
+
+    # Tasks in the task group don't depend on each other; they both become
+    # downstreams to task_start, and upstreams to task_end.
+    assert task_thing_1.upstream_task_ids == task_thing_2.upstream_task_ids == {"task_start"}
+    assert task_thing_1.downstream_task_ids == task_thing_2.downstream_task_ids == {"task_end"}
+
+
 def test_build_task_group_with_operators():
     """Tests DAG with Tasks created with *Operators and TaskGroup created with taskgroup decorator"""
 
@@ -731,7 +767,7 @@ def section_a(value):
         t_end = PythonOperator(task_id='task_end', python_callable=task_end, dag=dag)
         sec_1.set_downstream(t_end)
 
-    # Testing Tasks ing DAG
+    # Testing Tasks in DAG
     assert set(dag.task_group.children.keys()) == {'section_1', 'task_start', 'task_end'}
     assert set(dag.task_group.children['section_1'].children.keys()) == {
         'section_1.task_2',