This is an automated email from the ASF dual-hosted git repository. jedcunningham pushed a commit to branch v2-3-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit ecbae242c69314c83e97b3baa385f6887f76d9c0 Author: sanjayp <[email protected]> AuthorDate: Mon May 23 06:54:52 2022 +0530 Add __wrapped__ property to _TaskDecorator (#23830) Co-authored-by: Sanjay Pillai <sanjaypillai11 [at] gmail.com> (cherry picked from commit a71e4b789006b8f36cd993731a9fb7d5792fccc2) --- airflow/decorators/base.py | 10 +++++++++- tests/decorators/test_python.py | 17 +++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py index 2029f6c5ed..79277c7281 100644 --- a/airflow/decorators/base.py +++ b/airflow/decorators/base.py @@ -267,6 +267,10 @@ class _TaskDecorator(Generic[Function, OperatorSubclass]): op.doc_md = self.function.__doc__ return XComArg(op) + @property + def __wrapped__(self) -> Function: + return self.function + @cached_property def function_signature(self): return inspect.signature(self.function) @@ -495,6 +499,10 @@ class Task(Generic[Function]): function: Function + @property + def __wrapped__(self) -> Function: + ... + def expand(self, **kwargs: "Mappable") -> XComArg: ... @@ -527,7 +535,7 @@ def task_decorator_factory( **kwargs, ) -> TaskDecorator: """ - A factory that generates a wrapper that raps a function into an Airflow operator. + A factory that generates a wrapper that wraps a function into an Airflow operator. Accepts kwargs for operator kwarg. Can be reused in a single DAG. :param python_callable: Function to decorate diff --git a/tests/decorators/test_python.py b/tests/decorators/test_python.py index 3f8b44c464..63514fcf25 100644 --- a/tests/decorators/test_python.py +++ b/tests/decorators/test_python.py @@ -738,3 +738,20 @@ def test_mapped_render_template_fields(dag_maker, session): assert op.op_kwargs['arg1'] == "{{ ds }}" assert op.op_kwargs['arg2'] == "fn" + + +def test_task_decorator_has_wrapped_attr(): + """ + Test @task original underlying function is accessible + through the __wrapped__ attribute. + """ + + def org_test_func(): + pass + + decorated_test_func = task_decorator(org_test_func) + + assert hasattr( + decorated_test_func, '__wrapped__' + ), "decorated function does not have __wrapped__ attribute" + assert decorated_test_func.__wrapped__ is org_test_func, "__wrapped__ attr is not the original function"
