Repository: apache/airflow

Issue:
TaskFlow API `multiple_outputs` inferral causes import errors when using TYPE_CHECKING
### Apache Airflow version

2.5.1

### What happened

When using the TaskFlow API, I like to generally keep a good practice of adding type annotations in the TaskFlow functions so others reading the DAG and task code have better context around inputs/outputs, keep imports solely used for typing behind `typing.TYPE_CHECKING`, and utilize PEP 563 for forwarding annotation evaluations. Unfortunately, when using ~PEP 563 _and_ `TYPE_CHECKING`~ just TYPE_CHECKING, DAG import errors occur with a "NameError: <name> is not defined." exception.

### What you think should happen instead

Users should be free to use ~PEP 563 and~ `TYPE_CHECKING` when using the TaskFlow API and not hit DAG import errors along the way.

### How to reproduce

Using a straightforward use case of transforming a DataFrame, let's assume this toy example:

```py
from __future__ import annotations

from typing import TYPE_CHECKING, Any

from pendulum import datetime

from airflow.decorators import dag, task

if TYPE_CHECKING:
    from pandas import DataFrame


@dag(start_date=datetime(2023, 1, 1), schedule=None)
def multiple_outputs():
    @task()
    def transform(df: DataFrame) -> dict[str, Any]:
        ...

    transform()

multiple_outputs()
```
Add this DAG to your DAGS_FOLDER and the following import error should be observed:
<img width="641" alt="image" src="https://user-images.githubusercontent.com/48934154/217713685-ec29d5cc-4a48-4049-8dfa-56cbd76cddc3.png">


### Operating System

Debian GNU/Linux

### Versions of Apache Airflow Providers

apache-airflow-providers-amazon==6.2.0
apache-airflow-providers-apache-hive==5.1.1
apache-airflow-providers-apache-livy==3.2.0
apache-airflow-providers-celery==3.1.0
apache-airflow-providers-cncf-kubernetes==5.1.1
apache-airflow-providers-common-sql==1.3.3
apache-airflow-providers-databricks==4.0.0
apache-airflow-providers-dbt-cloud==2.3.1
apache-airflow-providers-elasticsearch==4.3.3
apache-airflow-providers-ftp==3.3.0
apache-airflow-providers-google==8.8.0
apache-airflow-providers-http==4.1.1
apache-airflow-providers-imap==3.1.1
apache-airflow-providers-microsoft-azure==5.1.0
apache-airflow-providers-postgres==5.4.0
apache-airflow-providers-redis==3.1.0
apache-airflow-providers-sftp==4.2.1
apache-airflow-providers-snowflake==4.0.2
apache-airflow-providers-sqlite==3.3.1
apache-airflow-providers-ssh==3.4.0
astronomer-providers==1.14.0

### Deployment

Astronomer

### Deployment details

OOTB local Airflow install with LocalExecutor built with the Astro CLI.

### Anything else

- This behavior/error was not observed using Airflow 2.4.3. 
- As a workaround, `multiple_outputs` can be explicitly set on the TaskFlow function to skip the inferral.

### Are you willing to submit 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)



Discussion:
Oh how I wish Lazy imports would be already implemented and available for use. https://peps.python.org/pep-0690/ All this __fuutre_annotations and TYPE_CHECKING dance would not be needed. 
Actually this looks like an issue even if you don't use PEP 563 and have stringified typing:
```python
# from __future__ import annotations

from typing import TYPE_CHECKING, Any

from pendulum import datetime

from airflow.decorators import dag, task

if TYPE_CHECKING:
    from pandas import DataFrame


@dag(start_date=datetime(2023, 1, 1), schedule=None)
def multiple_outputs():
    @task()
    def transform(df: "DataFrame" = None) -> dict[str, Any]:
        ...

    transform()

multiple_outputs()
```
Same `NameError` when parsing the DAG.
What's the full stack trace for the import error please?
I think there's two things we can do here: First off, we don't have to look at the annotations for anything but the return value (so a lazy arg annotation shouldn't matter as we can ignore it) and secondly is use `typing.get_type_hints()` to turn the annotation string to an actual type.
Oh, we're already using `get_type_hints()`


```
airflow/decorators/python.py:75: in python_task
    return task_decorator_factory(
airflow/decorators/base.py:572: in task_decorator_factory
    decorator = _TaskDecorator(
<attrs generated init airflow.decorators.base._TaskDecorator>:7: in __init__
    _setattr(self, 'multiple_outputs', __attr_factory_multiple_outputs(self))
airflow/decorators/base.py:299: in _infer_multiple_outputs
    return_type = typing_extensions.get_type_hints(self.function).get("return", Any)
../../../.virtualenvs/airflow/lib/python3.10/site-packages/typing_extensions.py:1298: in get_type_hints
    hint = typing.get_type_hints(
/usr/lib/python3.10/typing.py:1871: in get_type_hints
    value = _eval_type(value, globalns, localns)
/usr/lib/python3.10/typing.py:327: in _eval_type
    return t._evaluate(globalns, localns, recursive_guard)
/usr/lib/python3.10/typing.py:694: in _evaluate
    eval(self.__forward_code__, globalns, localns),
```
Fix coming.
Full stack trace:
```
[2023-02-09 15:48:34,330] {dagbag.py:343} ERROR - Failed to import: /usr/local/airflow/dags/mo_inferral.py
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/airflow/models/dagbag.py", line 339, in parse
    loader.exec_module(new_module)
  File "<frozen importlib._bootstrap_external>", line 850, in exec_module
  File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
  File "/usr/local/airflow/dags/mo_inferral.py", line 21, in <module>
    multiple_outputs()
  File "/usr/local/lib/python3.9/site-packages/airflow/models/dag.py", line 3555, in factory
    f(**f_kwargs)
  File "/usr/local/airflow/dags/mo_inferral.py", line 16, in multiple_outputs
    def transform(df: "DataFrame" = None) -> dict[str, Any]:
  File "/usr/local/lib/python3.9/site-packages/airflow/decorators/base.py", line 583, in decorator_factory
    return _TaskDecorator(
  File "<attrs generated init airflow.decorators.base._TaskDecorator>", line 8, in __init__
    _setattr('multiple_outputs', __attr_factory_multiple_outputs(self))
  File "/usr/local/lib/python3.9/site-packages/airflow/decorators/base.py", line 299, in _infer_multiple_outputs
    return_type = typing_extensions.get_type_hints(self.function).get("return", Any)
  File "/usr/local/lib/python3.9/site-packages/typing_extensions.py", line 929, in get_type_hints
    hint = typing.get_type_hints(
  File "/usr/local/lib/python3.9/typing.py", line 1497, in get_type_hints
    value = _eval_type(value, globalns, localns)
  File "/usr/local/lib/python3.9/typing.py", line 292, in _eval_type
    return t._evaluate(globalns, localns, recursive_guard)
  File "/usr/local/lib/python3.9/typing.py", line 554, in _evaluate
    eval(self.__forward_code__, globalns, localns),
  File "<string>", line 1, in <module>
NameError: name 'DataFrame' is not defined
```
> Fix coming.

@ash Would you prefer I close my draft PR? https://github.com/apache/airflow/pull/29436
> > Fix coming.
> 
> @ash Would you prefer I close my draft PR? #29436

Keep it open for now, my approach may not be realistic.

Resolving patch:
diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -18,6 +18,7 @@
 
 import inspect
 import re
+import warnings
 from itertools import chain
 from typing import (
     Any,
@@ -295,8 +296,25 @@ class _TaskDecorator(ExpandableFactory, Generic[FParams, FReturn, OperatorSubcla
 
     @multiple_outputs.default
     def _infer_multiple_outputs(self):
+        if "return" not in self.function.__annotations__:
+            # No return type annotation, nothing to infer
+            return False
+
         try:
-            return_type = typing_extensions.get_type_hints(self.function).get("return", Any)
+            # We only care about the return annotation, not anything about the parameters
+            def fake():
+                ...
+
+            fake.__annotations__ = {"return": self.function.__annotations__["return"]}
+
+            return_type = typing_extensions.get_type_hints(fake, self.function.__globals__).get("return", Any)
+        except NameError as e:
+            warnings.warn(
+                f"Cannot infer multiple_outputs for TaskFlow function {self.function.__name__!r} with forward"
+                f" type references that are not imported. (Error was {e})",
+                stacklevel=4,
+            )
+            return False
         except TypeError:  # Can't evaluate return type.
             return False
         ttype = getattr(return_type, "__origin__", return_type)


Test patch:
diff --git a/tests/decorators/test_python.py b/tests/decorators/test_python.py
--- a/tests/decorators/test_python.py
+++ b/tests/decorators/test_python.py
@@ -18,7 +18,7 @@
 import sys
 from collections import namedtuple
 from datetime import date, timedelta
-from typing import Dict, Tuple  # noqa: F401  # This is used by annotation tests.
+from typing import TYPE_CHECKING, Dict, Tuple
 
 import pytest
 
@@ -88,6 +88,44 @@ def identity_dict_with_decorator_call(x: int, y: int) -> resolve(annotation):
 
         assert identity_dict_with_decorator_call(5, 5).operator.multiple_outputs is True
 
+    def test_infer_multiple_outputs_forward_annotation(self):
+        if TYPE_CHECKING:
+
+            class FakeTypeCheckingOnlyClass:
+                ...
+
+            class UnresolveableName:
+                ...
+
+        @task_decorator
+        def t1(x: "FakeTypeCheckingOnlyClass", y: int) -> Dict[int, int]:  # type: ignore[empty-body]
+            ...
+
+        assert t1(5, 5).operator.multiple_outputs is True
+
+        @task_decorator
+        def t2(x: "FakeTypeCheckingOnlyClass", y: int) -> "Dict[int, int]":  # type: ignore[empty-body]
+            ...
+
+        assert t2(5, 5).operator.multiple_outputs is True
+
+        with pytest.warns(UserWarning, match="Cannot infer multiple_outputs.*t3") as recwarn:
+
+            @task_decorator
+            def t3(  # type: ignore[empty-body]
+                x: "FakeTypeCheckingOnlyClass",
+                y: int,
+            ) -> "UnresolveableName[int, int]":
+                ...
+
+            line = sys._getframe().f_lineno - 3
+
+        warn = recwarn[0]
+        assert warn.filename == __file__
+        assert warn.lineno == line
+
+        assert t3(5, 5).operator.multiple_outputs is False
+
     def test_infer_multiple_outputs_using_other_typing(self):
         @task_decorator
         def identity_tuple(x: int, y: int) -> Tuple[int, int]: