This is an automated email from the ASF dual-hosted git repository.
jscheffl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e11b91c8e0 Make sure `multiple_outputs` is inferred correctly even
when using `TypedDict` (#36652)
e11b91c8e0 is described below
commit e11b91c8e01b38023f209983a81aee23439a34a3
Author: Noam Cohen <[email protected]>
AuthorDate: Wed Jan 10 23:21:21 2024 +0200
Make sure `multiple_outputs` is inferred correctly even when using
`TypedDict` (#36652)
* Use `issubclass()` to check if return type is a dictionary
* Compare type to `typing.Mapping` instead of `typing.Dict`
* Add documentation
---
airflow/decorators/base.py | 3 +--
docs/apache-airflow/tutorial/taskflow.rst | 4 ++--
tests/decorators/test_python.py | 13 +++++++++++++
3 files changed, 16 insertions(+), 4 deletions(-)
diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
index d3ec556f05..119672dd42 100644
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -27,7 +27,6 @@ from typing import (
Callable,
ClassVar,
Collection,
- Dict,
Generic,
Iterator,
Mapping,
@@ -351,7 +350,7 @@ class _TaskDecorator(ExpandableFactory, Generic[FParams,
FReturn, OperatorSubcla
except TypeError: # Can't evaluate return type.
return False
ttype = getattr(return_type, "__origin__", return_type)
- return ttype is dict or ttype is Dict
+ return issubclass(ttype, Mapping)
def __attrs_post_init__(self):
if "self" in self.function_signature.parameters:
diff --git a/docs/apache-airflow/tutorial/taskflow.rst
b/docs/apache-airflow/tutorial/taskflow.rst
index 5d71576b59..6d5b3bed25 100644
--- a/docs/apache-airflow/tutorial/taskflow.rst
+++ b/docs/apache-airflow/tutorial/taskflow.rst
@@ -428,8 +428,8 @@ Tasks can also infer multiple outputs by using dict Python
typing.
def identity_dict(x: int, y: int) -> dict[str, int]:
return {"x": x, "y": y}
-By using the typing ``Dict`` for the function return type, the
``multiple_outputs`` parameter
-is automatically set to true.
+By using the typing ``dict``, or any other class that conforms to the
``typing.Mapping`` protocol,
+for the function return type, the ``multiple_outputs`` parameter is
automatically set to true.
Note, If you manually set the ``multiple_outputs`` parameter the inference is
disabled and
the parameter value is used.
diff --git a/tests/decorators/test_python.py b/tests/decorators/test_python.py
index 78203ca6ef..98aab562b8 100644
--- a/tests/decorators/test_python.py
+++ b/tests/decorators/test_python.py
@@ -97,6 +97,19 @@ class TestAirflowTaskDecorator(BasePythonTest):
assert identity_dict_with_decorator_call(5,
5).operator.multiple_outputs is True
+ @pytest.mark.skipif(sys.version_info < (3, 8), reason="PEP 589 is
implemented in Python 3.8")
+ def test_infer_multiple_outputs_typed_dict(self):
+ from typing import TypedDict
+
+ class TypeDictClass(TypedDict):
+ pass
+
+ @task_decorator
+ def t1() -> TypeDictClass:
+ return {}
+
+ assert t1().operator.multiple_outputs is True
+
def test_infer_multiple_outputs_forward_annotation(self):
if TYPE_CHECKING:
