Skip to content

Commit

Permalink
feat(sdk): Python components - Parse component input/output descripti…
Browse files Browse the repository at this point in the history
…ons 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
  • Loading branch information
munagekar committed Sep 20, 2020
1 parent 4e7877f commit 5613db0
Show file tree
Hide file tree
Showing 5 changed files with 35 additions and 8 deletions.
19 changes: 11 additions & 8 deletions sdk/python/kfp/components/_python_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,11 @@

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')


Expand Down Expand Up @@ -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 = []

Expand Down Expand Up @@ -318,6 +323,7 @@ def annotation_to_type_struct(annotation):
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
Expand All @@ -328,6 +334,7 @@ def annotation_to_type_struct(annotation):
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
Expand Down Expand Up @@ -387,14 +394,10 @@ def annotation_to_type_struct(annotation):
# 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,
Expand Down
21 changes: 21 additions & 0 deletions sdk/python/kfp/components_tests/test_python_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -454,6 +454,27 @@ def add_multiply_two_numbers(a: float = 3, b: float = 5) -> NamedTuple('DummyNam
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
Expand Down
1 change: 1 addition & 0 deletions sdk/python/requirements.in
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ PyYAML
# kfp.components
cloudpickle
strip-hints>=0.1.8
docstring-parser>=0.7.3

# kfp.dsl
jsonschema>=3.0.1
Expand Down
1 change: 1 addition & 0 deletions sdk/python/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions sdk/python/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
'click',
'Deprecated',
'strip-hints',
'docstring-parser>=0.7.3'
]

TESTS_REQUIRE = [
Expand Down

0 comments on commit 5613db0

Please sign in to comment.