Skip to content

Commit

Permalink
Use functools.update_wrapper on task decorator (apache#42079)
Browse files Browse the repository at this point in the history
  • Loading branch information
uranusjr authored Sep 8, 2024
1 parent 86959bf commit 8f665a5
Showing 1 changed file with 4 additions and 8 deletions.
12 changes: 4 additions & 8 deletions airflow/decorators/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import itertools
import textwrap
import warnings
from functools import cached_property
from functools import cached_property, update_wrapper
from typing import (
TYPE_CHECKING,
Any,
Expand Down Expand Up @@ -334,9 +334,8 @@ class _TaskDecorator(ExpandableFactory, Generic[FParams, FReturn, OperatorSubcla
is_teardown: bool = False
on_failure_fail_dagrun: bool = False

@property
def __doc__(self):
return self.function.__doc__
# This is set in __attrs_post_init__ by update_wrapper. Provided here for type hints.
__wrapped__: Callable[FParams, FReturn] = attr.ib(init=False)

@multiple_outputs.default
def _infer_multiple_outputs(self):
Expand Down Expand Up @@ -367,6 +366,7 @@ def __attrs_post_init__(self):
if "self" in self.function_signature.parameters:
raise TypeError(f"@{self.decorator_name} does not support methods")
self.kwargs.setdefault("task_id", self.function.__name__)
update_wrapper(self, self.function)

def __call__(self, *args: FParams.args, **kwargs: FParams.kwargs) -> XComArg:
if self.is_teardown:
Expand All @@ -390,10 +390,6 @@ def __call__(self, *args: FParams.args, **kwargs: FParams.kwargs) -> XComArg:
op.doc_md = self.function.__doc__
return XComArg(op)

@property
def __wrapped__(self) -> Callable[FParams, FReturn]:
return self.function

def _validate_arg_names(self, func: ValidationSource, kwargs: dict[str, Any]):
# Ensure that context variables are not shadowed.
context_keys_being_mapped = KNOWN_CONTEXT_KEYS.intersection(kwargs)
Expand Down

0 comments on commit 8f665a5

Please sign in to comment.