From 5ab368ac105968ccb79b2bcd32d3413f5e318a56 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Fri, 1 Mar 2019 14:51:18 -0800 Subject: [PATCH] Added support for default values to Lightweight python components (#890) --- sdk/python/kfp/components/_python_op.py | 1 + sdk/python/tests/components/test_python_op.py | 12 ++++++++++++ 2 files changed, 13 insertions(+) diff --git a/sdk/python/kfp/components/_python_op.py b/sdk/python/kfp/components/_python_op.py index 8b53e095d85..ce3d9c795d0 100644 --- a/sdk/python/kfp/components/_python_op.py +++ b/sdk/python/kfp/components/_python_op.py @@ -96,6 +96,7 @@ def annotation_to_type_struct(annotation): input_spec = InputSpec( name=parameter.name, type=type_struct, + default=str(parameter.default) if parameter.default is not inspect.Parameter.empty else None, ) inputs.append(input_spec) diff --git a/sdk/python/tests/components/test_python_op.py b/sdk/python/tests/components/test_python_op.py index 6ba97ddfab7..f82a0946d62 100644 --- a/sdk/python/tests/components/test_python_op.py +++ b/sdk/python/tests/components/test_python_op.py @@ -181,6 +181,18 @@ def add_two_numbers_decorated(a: float, b: float) -> float: self.assertEqual(component_spec.description.strip(), expected_description.strip()) self.assertEqual(component_spec.implementation.container.image, expected_image) + def test_saving_default_values(self): + from typing import NamedTuple + def add_multiply_two_numbers(a: float = 3, b: float = 5) -> NamedTuple('DummyName', [('sum', float), ('product', float)]): + '''Returns sum and product of two arguments''' + return (a + b, a * b) + + func = add_multiply_two_numbers + component_spec = comp._python_op._func_to_component_spec(func) + + self.assertEqual(component_spec.inputs[0].default, '3') + self.assertEqual(component_spec.inputs[1].default, '5') + def test_end_to_end_python_component_pipeline_compilation(self): import kfp.components as comp