fix(sdk): fix extract docstring in load component (#7921)

* fix extract docstring in load component

* implement review feedback
This commit is contained in:
Connor McCarthy 2022-06-24 12:12:35 -06:00 committed by GitHub
parent 281151c051
commit a020c9b01c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 5 additions and 5 deletions

View File

@ -69,8 +69,7 @@ CONFIG = {
},
'v1_components': {
'test_cases': [
# TODO: this component currently has a bug when extracting the component description from the command
# 'concat_placeholder_component',
'concat_placeholder_component',
# TODO: these three currently have placeholder bugs -- uncomment after fix
# 'if_placeholder_component',
# 'add_component',

View File

@ -674,7 +674,8 @@ class ComponentSpec(base_model.BaseModel):
for name, parameter_dict in all_outputs.items()
}
def extract_description_from_command(commands: List[str]) -> str:
def extract_description_from_command(
commands: List[str]) -> Union[str, None]:
for command in commands:
if isinstance(command, str) and 'import kfp' in command:
for node in ast.walk(ast.parse(command)):
@ -684,7 +685,7 @@ class ComponentSpec(base_model.BaseModel):
docstring = ast.get_docstring(node)
if docstring:
return docstring
return ''
return None
inputs = inputs_dict_from_components_dict(
pipeline_spec_dict['components'], raw_name)
@ -692,7 +693,7 @@ class ComponentSpec(base_model.BaseModel):
pipeline_spec_dict['components'], raw_name)
description = extract_description_from_command(
implementation.container.command) or None
implementation.container.command or [])
return ComponentSpec(
name=raw_name,
implementation=implementation,