Skip to content
Open
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
45 changes: 41 additions & 4 deletions task-sdk/src/airflow/sdk/bases/decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,6 +253,38 @@ def determine_kwargs(
return KeywordParameters.determine(func, args, kwargs).unpacking()


_TASK_DECORATOR_CALL_HINT = (
"This can happen when a @task-decorated function shadows another callable and the decorated task "
"object is called like a regular function. Rename the task function or call the original callable instead."
)


def _is_python_callable_already_executing(python_callable: Callable) -> bool:
target_code = getattr(python_callable, "__code__", None)
if target_code is None:
return False

frame = inspect.currentframe()
try:
while frame is not None:
if frame.f_code is target_code:
return True
frame = frame.f_back
return False
finally:
del frame


def _should_add_task_decorator_call_hint(
err: TypeError, python_callable: Callable, op_args: Collection[Any]
) -> bool:
return (
bool(op_args)
and "too many positional arguments" in str(err)
and _is_python_callable_already_executing(python_callable)
)


class DecoratedOperator(BaseOperator):
"""
Wraps a Python callable and captures args/kwargs when called for execution.
Expand Down Expand Up @@ -365,10 +397,15 @@ def __init__(
# check all the arguments we know are valid. Whether these are enough
# can only be known at execution time, when unmapping happens, and this
# is called without the _airflow_mapped_validation_only flag.
if kwargs.get("_airflow_mapped_validation_only"):
signature.bind_partial(*op_args, **op_kwargs)
else:
signature.bind(*op_args, **op_kwargs)
try:
if kwargs.get("_airflow_mapped_validation_only"):
signature.bind_partial(*op_args, **op_kwargs)
else:
signature.bind(*op_args, **op_kwargs)
except TypeError as err:
if _should_add_task_decorator_call_hint(err, python_callable, op_args):
raise TypeError(f"{err}. {_TASK_DECORATOR_CALL_HINT}") from err
raise

# Params in injected_for_ordering are semantically required even though they received a
# None default to satisfy Python's ordering constraint. Verify they are actually provided.
Expand Down
30 changes: 30 additions & 0 deletions task-sdk/tests/task_sdk/bases/test_decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,36 @@ def dummy_task(required_arg):
with pytest.raises(TypeError):
make_op(dummy_task)

def test_bind_validation_hints_for_accidental_task_decorator_call(self):
@task
def sleep():
sleep(3600)

with pytest.raises(
TypeError,
match="too many positional arguments.*@task-decorated function shadows another callable",
):
sleep.function()

def test_bind_validation_plain_arity_error_has_no_accidental_call_hint(self):
@task
def dummy_task(required_arg):
return required_arg

with pytest.raises(TypeError) as ctx:
dummy_task(1, 2)

assert "@task-decorated function shadows another callable" not in str(ctx.value)

def test_bind_validation_missing_required_args_has_no_accidental_call_hint(self):
def dummy_task(required_arg):
return required_arg

with pytest.raises(TypeError) as ctx:
make_op(dummy_task)

assert "@task-decorated function shadows another callable" not in str(ctx.value)

def test_variadic_and_keyword_only_params_are_not_assigned_defaults(self):
"""Construction succeeds when variadic and keyword-only params are present."""

Expand Down