SDK - Python components - Fixed handling multiline decorators

sdk/python/kfp/components/_python_op.py

@@ -190,26 +190,29 @@ def _capture_function_code_using_cloudpickle(func, modules_to_capture: List[str]
 
 
 def _capture_function_code_using_source_copy(func) -> str:	
-    import inspect	
+    import textwrap
 
-    #Source code can include decorators line @python_op. Remove them
-    (func_code_lines, _) = inspect.getsourcelines(func)
-    while func_code_lines[0].lstrip().startswith('@'): #decorator
-        del func_code_lines[0]
+    func_code = inspect.getsource(func)
 
     #Function might be defined in some indented scope (e.g. in another function).
     #We need to handle this and properly dedent the function source code
-    first_line = func_code_lines[0]
-    indent = len(first_line) - len(first_line.lstrip())
-    func_code_lines = [line[indent:] for line in func_code_lines]
+    func_code = textwrap.dedent(func_code)
+    func_code_lines = func_code.split('\n')
+
+    # Removing possible decorators (can be multiline) until the function definition is found
+    while func_code_lines and not func_code_lines[0].startswith('def '):
+        del func_code_lines[0]
+
+    if not func_code_lines:
+        raise ValueError('Failed to dedent and clean up the source of function "{}". It is probably not properly indented.'.format(func.__name__))
 
     #TODO: Add support for copying the NamedTuple subclass declaration code
     #Adding NamedTuple import if needed
     if hasattr(inspect.signature(func).return_annotation, '_fields'): #NamedTuple
-        func_code_lines.insert(0, '\n')
-        func_code_lines.insert(0, 'from typing import NamedTuple\n')
+        func_code_lines.insert(0, '')
+        func_code_lines.insert(0, 'from typing import NamedTuple')
 
-    return ''.join(func_code_lines) #Lines retain their \n endings
+    return '\n'.join(func_code_lines)
 
 
 def _extract_component_interface(func) -> ComponentSpec:

sdk/python/tests/components/test_python_op.py

@@ -17,7 +17,7 @@
 import unittest
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Sequence
 
 import kfp
 import kfp.components as comp
@@ -122,9 +122,11 @@ def helper_test_component_against_func_using_local_call(self, func: Callable, op
         # ! This function cannot be used when component has output types that use custom serialization since it will compare non-serialized function outputs with serialized component outputs.
         # Evaluating the function to get the expected output values
         expected_output_values_list = func(**arguments)
+        if not isinstance(expected_output_values_list, Sequence) or isinstance(expected_output_values_list, str):
+            expected_output_values_list = [str(expected_output_values_list)]
         expected_output_values_list = [str(value) for value in expected_output_values_list]
 
-        output_names = [output.name for output in op.outputs]
+        output_names = [output.name for output in op.component_spec.outputs]
         from kfp.components._naming import generate_unique_name_conversion_table, _sanitize_python_function_name
         output_name_to_pythonic = generate_unique_name_conversion_table(output_names, _sanitize_python_function_name)
         pythonic_output_names = [output_name_to_pythonic[name] for name in output_names]
@@ -401,8 +403,15 @@ def test_python_component_decorator(self):
         expected_description = 'Sum component description'
         expected_image = 'org/image'
 
-        @python_component(name=expected_name, description=expected_description, base_image=expected_image)
-        def add_two_numbers_decorated(a: float, b: float) -> float:
+        @python_component(
+            name=expected_name,
+            description=expected_description,
+            base_image=expected_image
+        )
+        def add_two_numbers_decorated(
+            a: float,
+            b: float,
+        ) -> float:
             '''Returns sum of two arguments'''
             return a + b
 
@@ -412,6 +421,11 @@ 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)
 
+        func = add_two_numbers_decorated
+        op = comp.func_to_container_op(func)
+
+        self.helper_test_component_against_func_using_local_call(func, op, arguments={'a': 3, 'b': 5.0})
+
     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)]):