<h3 dir="auto">Apache Airflow version</h3>
<p dir="auto">2.5.3</p>
<h3 dir="auto">What happened</h3>
<p dir="auto">Pleas consider the following scenario</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from my_package import MyClass

@task
def foo(a: int, b: str, c: &quot;MyClass&quot;): 
    pass"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">typing</span> <span class="pl-k">import</span> <span class="pl-v">TYPE_CHECKING</span>
<span class="pl-k">if</span> <span class="pl-v">TYPE_CHECKING</span>:
    <span class="pl-k">from</span> <span class="pl-s1">my_package</span> <span class="pl-k">import</span> <span class="pl-v">MyClass</span>

<span class="pl-en">@<span class="pl-s1">task</span></span>
<span class="pl-k">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">a</span>: <span class="pl-s1">int</span>, <span class="pl-s1">b</span>: <span class="pl-s1">str</span>, <span class="pl-s1">c</span>: <span class="pl-s">"MyClass"</span>): 
    <span class="pl-k">pass</span></pre></div>
<p dir="auto">In this case, the DAG parsing will fail with a NameError from the <code class="notranslate">_infer_multiple_outputs</code> method of <code class="notranslate">_TaskDecorator</code> class in <code class="notranslate">airflow/decorators/base.py</code>.</p>
<p dir="auto">Please note that this is a common pattern when managing multiple tasks with varying dependencies and docker containers, and <code class="notranslate">my_package</code> might not be available in some of them.</p>
<h3 dir="auto">What you think should happen instead</h3>
<p dir="auto">The problem relies in the following method using typing_extensions.get_type_hints:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="@multiple_outputs.default
def _infer_multiple_outputs(self):
    try:
        return_type = typing_extensions.get_type_hints(self.function).get(&quot;return&quot;, Any)
    except TypeError:  # Can't evaluate return type.
        return False
    ttype = getattr(return_type, &quot;__origin__&quot;, return_type)
    return ttype == dict or ttype == Dict"><pre class="notranslate"><span class="pl-en">@<span class="pl-s1">multiple_outputs</span>.<span class="pl-s1">default</span></span>
<span class="pl-k">def</span> <span class="pl-en">_infer_multiple_outputs</span>(<span class="pl-s1">self</span>):
    <span class="pl-k">try</span>:
        <span class="pl-s1">return_type</span> <span class="pl-c1">=</span> <span class="pl-s1">typing_extensions</span>.<span class="pl-en">get_type_hints</span>(<span class="pl-s1">self</span>.<span class="pl-s1">function</span>).<span class="pl-en">get</span>(<span class="pl-s">"return"</span>, <span class="pl-v">Any</span>)
    <span class="pl-k">except</span> <span class="pl-v">TypeError</span>:  <span class="pl-c"># Can't evaluate return type.</span>
        <span class="pl-k">return</span> <span class="pl-c1">False</span>
    <span class="pl-s1">ttype</span> <span class="pl-c1">=</span> <span class="pl-en">getattr</span>(<span class="pl-s1">return_type</span>, <span class="pl-s">"__origin__"</span>, <span class="pl-s1">return_type</span>)
    <span class="pl-k">return</span> <span class="pl-s1">ttype</span> <span class="pl-c1">==</span> <span class="pl-s1">dict</span> <span class="pl-c1">or</span> <span class="pl-s1">ttype</span> <span class="pl-c1">==</span> <span class="pl-v">Dict</span></pre></div>
<p dir="auto">which will result in the following error at DAG parse time:<br>
<code class="notranslate">E   NameError: name 'MyClass' is not defined</code></p>
<p dir="auto">To fix this, the method could just use the <code class="notranslate">__annotations__</code> attribute of <code class="notranslate">self.function</code> instead which does not complain about the existence of the annotated types, like so:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="@multiple_outputs.default
def _infer_multiple_outputs(self):
    try:
        return_type = self.function.__annotations__.get(&quot;return&quot;, Any)
    except TypeError:  # Can't evaluate return type.
        return False
    ttype = getattr(return_type, &quot;__origin__&quot;, return_type)
    return ttype == dict or ttype == Dict"><pre class="notranslate"><span class="pl-en">@<span class="pl-s1">multiple_outputs</span>.<span class="pl-s1">default</span></span>
<span class="pl-k">def</span> <span class="pl-en">_infer_multiple_outputs</span>(<span class="pl-s1">self</span>):
    <span class="pl-k">try</span>:
        <span class="pl-s1">return_type</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">function</span>.<span class="pl-s1">__annotations__</span>.<span class="pl-en">get</span>(<span class="pl-s">"return"</span>, <span class="pl-v">Any</span>)
    <span class="pl-k">except</span> <span class="pl-v">TypeError</span>:  <span class="pl-c"># Can't evaluate return type.</span>
        <span class="pl-k">return</span> <span class="pl-c1">False</span>
    <span class="pl-s1">ttype</span> <span class="pl-c1">=</span> <span class="pl-en">getattr</span>(<span class="pl-s1">return_type</span>, <span class="pl-s">"__origin__"</span>, <span class="pl-s1">return_type</span>)
    <span class="pl-k">return</span> <span class="pl-s1">ttype</span> <span class="pl-c1">==</span> <span class="pl-s1">dict</span> <span class="pl-c1">or</span> <span class="pl-s1">ttype</span> <span class="pl-c1">==</span> <span class="pl-v">Dict</span></pre></div>
<p dir="auto">This should work fine since we are interested in the return type annotation anyway.</p>
<p dir="auto">One quick fix for the programmers having this issue could simply be to add an <code class="notranslate">else</code> clause to their <code class="notranslate">if TYPE_CHECKING</code> statement, like so:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="if TYPE_CHECKING:
   from my_package import MyClass
else:
   from typing import Any as MyClass"><pre class="notranslate"><span class="pl-k">if</span> <span class="pl-v">TYPE_CHECKING</span>:
   <span class="pl-k">from</span> <span class="pl-s1">my_package</span> <span class="pl-k">import</span> <span class="pl-v">MyClass</span>
<span class="pl-k">else</span>:
   <span class="pl-k">from</span> <span class="pl-s1">typing</span> <span class="pl-k">import</span> <span class="pl-v">Any</span> <span class="pl-k">as</span> <span class="pl-v">MyClass</span></pre></div>
<h3 dir="auto">How to reproduce</h3>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from airflow.annotations import task
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from my_package import MyClass

@task
def foo(a: int, b: str, c: &quot;MyClass&quot;): 
    pass"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">airflow</span>.<span class="pl-s1">annotations</span> <span class="pl-k">import</span> <span class="pl-s1">task</span>
<span class="pl-k">from</span> <span class="pl-s1">typing</span> <span class="pl-k">import</span> <span class="pl-v">TYPE_CHECKING</span>
<span class="pl-k">if</span> <span class="pl-v">TYPE_CHECKING</span>:
    <span class="pl-k">from</span> <span class="pl-s1">my_package</span> <span class="pl-k">import</span> <span class="pl-v">MyClass</span>

<span class="pl-en">@<span class="pl-s1">task</span></span>
<span class="pl-k">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">a</span>: <span class="pl-s1">int</span>, <span class="pl-s1">b</span>: <span class="pl-s1">str</span>, <span class="pl-s1">c</span>: <span class="pl-s">"MyClass"</span>): 
    <span class="pl-k">pass</span></pre></div>
<h3 dir="auto">Operating System</h3>
<p dir="auto">Linux Mint 21</p>
<h3 dir="auto">Versions of Apache Airflow Providers</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">Deployment</h3>
<p dir="auto">Official Apache Airflow Helm Chart</p>
<h3 dir="auto">Deployment details</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">Anything else</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">Are you willing to submit PR?</h3>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Yes I am willing to submit a PR!</li>
</ul>
<h3 dir="auto">Code of Conduct</h3>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> I agree to follow this project's <a href="https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md">Code of Conduct</a></li>
</ul>
<h3 dir="auto">Apache Airflow version</h3>
<p dir="auto">2.5.1</p>
<h3 dir="auto">What happened</h3>
<p dir="auto">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 <code class="notranslate">typing.TYPE_CHECKING</code>, and utilize PEP 563 for forwarding annotation evaluations. Unfortunately, when using <del>PEP 563 <em>and</em> <code class="notranslate">TYPE_CHECKING</code></del> just TYPE_CHECKING, DAG import errors occur with a "NameError:  is not defined." exception.</p>
<h3 dir="auto">What you think should happen instead</h3>
<p dir="auto">Users should be free to use <del>PEP 563 and</del> <code class="notranslate">TYPE_CHECKING</code> when using the TaskFlow API and not hit DAG import errors along the way.</p>
<h3 dir="auto">How to reproduce</h3>
<p dir="auto">Using a straightforward use case of transforming a DataFrame, let's assume this toy example:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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) -&gt; dict[str, Any]:
        ...

    transform()

multiple_outputs()"><pre class="notranslate"><span class="pl-k">from</span> __future__ <span class="pl-k">import</span> <span class="pl-s1">annotations</span>

<span class="pl-k">from</span> <span class="pl-s1">typing</span> <span class="pl-k">import</span> <span class="pl-v">TYPE_CHECKING</span>, <span class="pl-v">Any</span>

<span class="pl-k">from</span> <span class="pl-s1">pendulum</span> <span class="pl-k">import</span> <span class="pl-s1">datetime</span>

<span class="pl-k">from</span> <span class="pl-s1">airflow</span>.<span class="pl-s1">decorators</span> <span class="pl-k">import</span> <span class="pl-s1">dag</span>, <span class="pl-s1">task</span>

<span class="pl-k">if</span> <span class="pl-v">TYPE_CHECKING</span>:
    <span class="pl-k">from</span> <span class="pl-s1">pandas</span> <span class="pl-k">import</span> <span class="pl-v">DataFrame</span>


<span class="pl-en">@<span class="pl-en">dag</span>(<span class="pl-s1">start_date</span><span class="pl-c1">=</span><span class="pl-en">datetime</span>(<span class="pl-c1">2023</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>), <span class="pl-s1">schedule</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)</span>
<span class="pl-k">def</span> <span class="pl-en">multiple_outputs</span>():
    <span class="pl-en">@<span class="pl-en">task</span>()</span>
    <span class="pl-k">def</span> <span class="pl-en">transform</span>(<span class="pl-s1">df</span>: <span class="pl-v">DataFrame</span>) <span class="pl-c1">-&gt;</span> <span class="pl-s1">dict</span>[<span class="pl-s1">str</span>, <span class="pl-v">Any</span>]:
        ...

    <span class="pl-en">transform</span>()

<span class="pl-en">multiple_outputs</span>()</pre></div>
<p dir="auto">Add this DAG to your DAGS_FOLDER and the following import error should be observed:<br>
<a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/48934154/217713685-ec29d5cc-4a48-4049-8dfa-56cbd76cddc3.png"><img width="641" alt="image" src="https://user-images.githubusercontent.com/48934154/217713685-ec29d5cc-4a48-4049-8dfa-56cbd76cddc3.png" style="max-width: 100%;"></a></p>
<h3 dir="auto">Operating System</h3>
<p dir="auto">Debian GNU/Linux</p>
<h3 dir="auto">Versions of Apache Airflow Providers</h3>
<p dir="auto">apache-airflow-providers-amazon==6.2.0<br>
apache-airflow-providers-apache-hive==5.1.1<br>
apache-airflow-providers-apache-livy==3.2.0<br>
apache-airflow-providers-celery==3.1.0<br>
apache-airflow-providers-cncf-kubernetes==5.1.1<br>
apache-airflow-providers-common-sql==1.3.3<br>
apache-airflow-providers-databricks==4.0.0<br>
apache-airflow-providers-dbt-cloud==2.3.1<br>
apache-airflow-providers-elasticsearch==4.3.3<br>
apache-airflow-providers-ftp==3.3.0<br>
apache-airflow-providers-google==8.8.0<br>
apache-airflow-providers-http==4.1.1<br>
apache-airflow-providers-imap==3.1.1<br>
apache-airflow-providers-microsoft-azure==5.1.0<br>
apache-airflow-providers-postgres==5.4.0<br>
apache-airflow-providers-redis==3.1.0<br>
apache-airflow-providers-sftp==4.2.1<br>
apache-airflow-providers-snowflake==4.0.2<br>
apache-airflow-providers-sqlite==3.3.1<br>
apache-airflow-providers-ssh==3.4.0<br>
astronomer-providers==1.14.0</p>
<h3 dir="auto">Deployment</h3>
<p dir="auto">Astronomer</p>
<h3 dir="auto">Deployment details</h3>
<p dir="auto">OOTB local Airflow install with LocalExecutor built with the Astro CLI.</p>
<h3 dir="auto">Anything else</h3>
<ul dir="auto">
<li>This behavior/error was not observed using Airflow 2.4.3.</li>
<li>As a workaround, <code class="notranslate">multiple_outputs</code> can be explicitly set on the TaskFlow function to skip the inferral.</li>
</ul>
<h3 dir="auto">Are you willing to submit PR?</h3>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Yes I am willing to submit a PR!</li>
</ul>
<h3 dir="auto">Code of Conduct</h3>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> I agree to follow this project's <a href="https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md">Code of Conduct</a></li>
</ul>