feat(sdk): Python components - Parse component input/output descriptions from the function docstring (#4512)

* cleanup imports

* add description to inputs and outputs

* update requirements

* add test

* improve component description

* update tests

* review changes: fix lint and requirements

* upgrade docstring-parser

sdk/python/kfp/components/_python_op.py

@@ -33,10 +33,11 @@ from .structures import *
 
 import inspect
 from pathlib import Path
-import typing
-from typing import Callable, Generic, List, TypeVar, Union
+from typing import Callable, List, TypeVar
 import warnings
 
+import docstring_parser
+
 T = TypeVar('T')
 
 
@@ -258,11 +259,15 @@ def _capture_function_code_using_source_copy(func) -> str:
     return func_code
 
 
-def _extract_component_interface(func) -> ComponentSpec:
+def _extract_component_interface(func: Callable) -> ComponentSpec:
     single_output_name_const = 'Output'
 
     signature = inspect.signature(func)
     parameters = list(signature.parameters.values())
+
+    parsed_docstring = docstring_parser.parse(inspect.getdoc(func))
+    doc_dict = {p.arg_name: p.description for p in parsed_docstring.params}
+
     inputs = []
     outputs = []
 
@@ -318,6 +323,7 @@ def _extract_component_interface(func) -> ComponentSpec:
             output_spec = OutputSpec(
                 name=io_name,
                 type=type_struct,
+                description=doc_dict.get(parameter.name)
             )
             output_spec._passing_style = passing_style
             output_spec._parameter_name = parameter.name
@@ -328,6 +334,7 @@ def _extract_component_interface(func) -> ComponentSpec:
             input_spec = InputSpec(
                 name=io_name,
                 type=type_struct,
+                description=doc_dict.get(parameter.name)
             )
             if parameter.default is not inspect.Parameter.empty:
                 input_spec.optional = True
@@ -387,14 +394,10 @@ def _extract_component_interface(func) -> ComponentSpec:
     # The name can be overridden by setting setting func.__name__ attribute (of the legacy func._component_human_name attribute).
     # The description can be overridden by setting the func.__doc__ attribute (or the legacy func._component_description attribute).
     component_name = getattr(func, '_component_human_name', None) or _python_function_name_to_component_name(func.__name__)
-    description = getattr(func, '_component_description', None) or func.__doc__
+    description = getattr(func, '_component_description', None) or parsed_docstring.short_description
     if description:
         description = description.strip()
 
-    # TODO: Parse input/output descriptions from the function docstring. See:
-    # https://github.com/rr-/docstring_parser
-    # https://github.com/terrencepreilly/darglint/blob/master/darglint/parse.py
-
     component_spec = ComponentSpec(
         name=component_name,
         description=description,

sdk/python/kfp/components_tests/test_python_op.py

@@ -454,6 +454,27 @@ class PythonOpTestCase(unittest.TestCase):
         self.assertEqual(component_spec.inputs[0].default, '3')
         self.assertEqual(component_spec.inputs[1].default, '5')
 
+    def test_handling_of_descriptions(self):
+
+        def pipeline(
+                env_var: str,
+                secret_name: str,
+                secret_key: str = None
+        ) -> None:
+            """
+            Pipeline to Demonstrate Usage of Secret
+
+            Args:
+                env_var: Name of the variable inside the Pod
+                secret_name: Name of the Secret in the namespace
+            """
+
+        component_spec = comp._python_op._func_to_component_spec(pipeline)
+        self.assertEqual(component_spec.description, 'Pipeline to Demonstrate Usage of Secret')
+        self.assertEqual(component_spec.inputs[0].description, 'Name of the variable inside the Pod')
+        self.assertEqual(component_spec.inputs[1].description, 'Name of the Secret in the namespace')
+        self.assertIsNone(component_spec.inputs[2].description)
+
     def test_handling_default_value_of_none(self):
         def assert_is_none(arg=None):
             assert arg is None

sdk/python/requirements.in

@@ -4,6 +4,7 @@ PyYAML
 # kfp.components
 cloudpickle
 strip-hints>=0.1.8
+docstring-parser>=0.7.3
 
 # kfp.dsl
 jsonschema>=3.0.1

sdk/python/requirements.txt

@@ -11,6 +11,7 @@ chardet==3.0.4            # via requests
 click==7.1.1              # via -r requirements.in
 cloudpickle==1.3.0        # via -r requirements.in
 deprecated==1.2.7         # via -r requirements.in
+docstring-parser==0.7.3   # via -r requirements.in
 google-api-core==1.16.0   # via google-cloud-core
 google-auth==1.11.3       # via -r requirements.in, google-api-core, google-cloud-storage, kubernetes
 google-cloud-core==1.3.0  # via google-cloud-storage

sdk/python/setup.py

@@ -37,6 +37,7 @@ REQUIRES = [
     'click',
     'Deprecated',
     'strip-hints',
+    'docstring-parser>=0.7.3'
 ]
 
 TESTS_REQUIRE = [