Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion airflow/operators/python.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)
Comment thread
ashb marked this conversation as resolved.
Outdated

def wrapper(f: T):
"""
Expand Down
16 changes: 16 additions & 0 deletions docs/apache-airflow/tutorial_taskflow_api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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?
------------

Expand Down
75 changes: 74 additions & 1 deletion tests/operators/test_python.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Comment thread
ashb marked this conversation as resolved.
Outdated
def test_fails_bad_signature(self):
"""Tests that @task will fail if signature is not binding."""

Expand Down