From 046f519153a2c8a8e0b2d2a84dbe3536eea525b3 Mon Sep 17 00:00:00 2001
From: Anderson Reyes <anderson@pymetrics.com>
Date: Fri, 13 Nov 2020 17:13:49 -0500
Subject: [PATCH] infer multiple outputs from dict annotations

---
 airflow/operators/python.py                   |  8 +-
 docs/apache-airflow/tutorial_taskflow_api.rst | 16 ++++
 tests/operators/test_python.py                | 75 ++++++++++++++++++-
 3 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/airflow/operators/python.py b/airflow/operators/python.py
index bb37c98dafa88..ec2aabe10a140 100644
--- a/airflow/operators/python.py
+++ b/airflow/operators/python.py
@@ -254,7 +254,7 @@ def execute(self, context: Dict):
 
 
 def task(
-    python_callable: Optional[Callable] = None, multiple_outputs: bool = False, **kwargs
+    python_callable: Optional[Callable] = None, multiple_outputs: Optional[bool] = None, **kwargs
 ) -> Callable[[T], T]:
     """
     Python operator decorator. Wraps a function into an Airflow operator.
@@ -269,6 +269,12 @@ def task(
     :type multiple_outputs: bool
 
     """
+    # try to infer from  type annotation
+    if python_callable and multiple_outputs is None:
+        sig = signature(python_callable).return_annotation
+        ttype = getattr(sig, "__origin__", None)
+
+        multiple_outputs = sig != inspect.Signature.empty and ttype in (dict, Dict)
 
     def wrapper(f: T):
         """
diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst
index 8154bd87d9ee9..6c51f1694585c 100644
--- a/docs/apache-airflow/tutorial_taskflow_api.rst
+++ b/docs/apache-airflow/tutorial_taskflow_api.rst
@@ -155,6 +155,22 @@ the dependencies as shown below.
     :end-before: [END main_flow]
 
 
+Multiple outputs inference
+--------------------------
+Tasks can also infer multiple outputs by using dict python typing.
+
+.. code-block:: python
+
+    @task
+    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.
+
+Note, If you manually set the ``multiple_outputs`` parameter the inference is disabled and
+the parameter value is used.
+
 What's Next?
 ------------
 
diff --git a/tests/operators/test_python.py b/tests/operators/test_python.py
index 95c6b18b53dfe..cd356335b4b11 100644
--- a/tests/operators/test_python.py
+++ b/tests/operators/test_python.py
@@ -22,7 +22,7 @@
 from collections import namedtuple
 from datetime import date, datetime, timedelta
 from subprocess import CalledProcessError
-from typing import List
+from typing import Dict, List, Tuple
 
 import funcsigs
 import pytest
@@ -329,6 +329,79 @@ def test_python_operator_python_callable_is_callable(self):
         with pytest.raises(AirflowException):
             task_decorator(not_callable, dag=self.dag)
 
+    def test_infer_multiple_outputs_using_typing(self):
+        @task_decorator
+        def identity_dict(x: int, y: int) -> Dict[str, int]:
+            return {"x": x, "y": y}
+
+        assert identity_dict(5, 5).operator.multiple_outputs is True  # pylint: disable=maybe-no-member
+
+        @task_decorator
+        def identity_tuple(x: int, y: int) -> Tuple[int, int]:
+            return x, y
+
+        assert identity_tuple(5, 5).operator.multiple_outputs is False  # pylint: disable=maybe-no-member
+
+        @task_decorator
+        def identity_int(x: int) -> int:
+            return x
+
+        assert identity_int(5).operator.multiple_outputs is False  # pylint: disable=maybe-no-member
+
+        @task_decorator
+        def identity_notyping(x: int):
+            return x
+
+        assert identity_notyping(5).operator.multiple_outputs is False  # pylint: disable=maybe-no-member
+
+    def test_manual_multiple_outputs_false_with_typings(self):
+        @task_decorator(multiple_outputs=False)
+        def identity2(x: int, y: int) -> Dict[int, int]:
+            return (x, y)
+
+        with self.dag:
+            res = identity2(8, 4)
+
+        dr = self.dag.create_dagrun(
+            run_id=DagRunType.MANUAL.value,
+            start_date=timezone.utcnow(),
+            execution_date=DEFAULT_DATE,
+            state=State.RUNNING,
+        )
+
+        res.operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)  # pylint: disable=maybe-no-member
+
+        ti = dr.get_task_instances()[0]
+
+        assert res.operator.multiple_outputs is False  # pylint: disable=maybe-no-member
+        assert ti.xcom_pull() == [8, 4]  # pylint: disable=maybe-no-member
+        assert ti.xcom_pull(key="return_value_0") is None
+        assert ti.xcom_pull(key="return_value_1") is None
+
+    def test_multiple_outputs_ignore_typing(self):
+        @task_decorator
+        def identity_tuple(x: int, y: int) -> Tuple[int, int]:
+            return x, y
+
+        with self.dag:
+            ident = identity_tuple(35, 36)
+
+        dr = self.dag.create_dagrun(
+            run_id=DagRunType.MANUAL.value,
+            start_date=timezone.utcnow(),
+            execution_date=DEFAULT_DATE,
+            state=State.RUNNING,
+        )
+
+        ident.operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)  # pylint: disable=maybe-no-member
+
+        ti = dr.get_task_instances()[0]
+
+        assert not ident.operator.multiple_outputs  # pylint: disable=maybe-no-member
+        assert ti.xcom_pull() == [35, 36]
+        assert ti.xcom_pull(key="return_value_0") is None
+        assert ti.xcom_pull(key="return_value_1") is None
+
     def test_fails_bad_signature(self):
         """Tests that @task will fail if signature is not binding."""