Set doc_md when using task decorator and function has __doc__ (#11598)

airflow/example_dags/tutorial_decorated_etl_dag.py

@@ -62,20 +62,25 @@ with DAG(
     # [START extract]
     @dag.task()
     def extract():
+        """
+        #### Extract task
+        A simple Extract task to get data ready for the rest of the data pipeline.
+        In this case, getting data is simulated by reading from a hardcoded JSON string.
+        """
         data_string = u'{"1001": 301.27, "1002": 433.21, "1003": 502.22}'
 
         order_data_dict = json.loads(data_string)
         return order_data_dict
     # [END extract]
-    extract.doc_md = """\
-#### Extract task
-A simple Extract task to get data ready for the rest of the data pipeline.
-In this case, getting data is simulated by reading from a hardcoded JSON string.
-"""
 
     # [START transform]
     @dag.task(multiple_outputs=True)
     def transform(order_data_dict: dict):
+        """
+        #### Transform task
+        A simple Transform task which takes in the collection of order data and computes
+        the total order value.
+        """
         total_order_value = 0
 
         for value in order_data_dict.values():
@@ -83,23 +88,18 @@ In this case, getting data is simulated by reading from a hardcoded JSON string.
 
         return {"total_order_value": total_order_value}
     # [END transform]
-    transform.doc_md = """\
-#### Transform task
-A simple Transform task which takes in the collection of order data and computes
-the total order value.
-"""
 
     # [START load]
     @dag.task()
     def load(total_order_value: float):
+        """
+        #### Load task
+        A simple Load task which takes in the result of the Transform task and instead of
+        saving it to end user review, just prints it out.
+        """
 
         print("Total order value is: %.2f" % total_order_value)
     # [END load]
-    load.doc_md = """\
-#### Load task
-A simple Load task which takes in the result of the Transform task and instead of
-saving it to end user review, just prints it out.
-"""
 
     # [START main_flow]
     order_data = extract()

airflow/operators/python.py

@@ -294,6 +294,8 @@ def task(
         def factory(*args, **f_kwargs):
             op = _PythonDecoratedOperator(python_callable=f, op_args=args, op_kwargs=f_kwargs,
                                           multiple_outputs=multiple_outputs, **kwargs)
+            if f.__doc__:
+                op.doc_md = f.__doc__
             return XComArg(op)
         return cast(T, factory)
     if callable(python_callable):

tests/operators/test_python.py

@@ -628,6 +628,22 @@ class TestAirflowTaskDecorator(TestPythonBase):
 
         assert 'add_2' in self.dag.task_ids
 
+    def test_task_documentation(self):
+        """Tests that task_decorator loads doc_md from function doc"""
+
+        @task_decorator
+        def add_2(number: int):
+            """
+            Adds 2 to number.
+            """
+            return number + 2
+
+        test_number = 10
+        with self.dag:
+            ret = add_2(test_number)
+
+        assert ret.operator.doc_md.strip(), "Adds 2 to number."  # pylint: disable=maybe-no-member
+
 
 class TestBranchOperator(unittest.TestCase):
     @classmethod