[airflow] branch main updated: Add __wrapped__ property to _TaskDecorator (#23830)

[uranusjr]

This is an automated email from the ASF dual-hosted git repository.

uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new a71e4b7890 Add __wrapped__ property to _TaskDecorator (#23830)
a71e4b7890 is described below

commit a71e4b789006b8f36cd993731a9fb7d5792fccc2
Author: sanjayp <sanjaypilla...@gmail.com>
AuthorDate: Mon May 23 06:54:52 2022 +0530

    Add __wrapped__ property to _TaskDecorator (#23830)
    
    Co-authored-by: Sanjay Pillai <sanjaypillai11 [at] gmail.com>
---
 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"