SDK/Components/Python - Made the typing.NamedTuple import optional (#717

)

Now it's only imported if the return type is NamedTuple.

sdk/python/kfp/components/_python_op.py

@@ -126,6 +126,12 @@ def annotation_to_type_struct(annotation):
 
     func_name=func.__name__
 
+    #TODO: Add support for copying the NamedTuple subclass declaration code
+    #Adding NamedTuple import if needed
+    func_type_declarations_code = ""
+    if hasattr(return_ann, '_fields'): #NamedTuple
+        func_type_declarations_code = func_type_declarations_code + '\n' + 'from typing import NamedTuple'
+
     #Source code can include decorators line @python_op. Remove them
     (func_code_lines, _) = inspect.getsourcelines(func) 
     while func_code_lines[0].lstrip().startswith('@'): #decorator
@@ -157,10 +163,10 @@ def annotation_to_type_struct(annotation):
 
     full_source = \
 '''\
-from typing import NamedTuple
-
 {extra_code}
 
+{func_type_declarations_code}
+
 {func_code}
 
 import sys
@@ -185,6 +191,7 @@ def annotation_to_type_struct(annotation):
 '''.format(
         func_name=func_name,
         func_code=func_code,
+        func_type_declarations_code=func_type_declarations_code,
         extra_code=extra_code,
         input_args_parsing_code='\n'.join(input_args_parsing_code_lines),
         output_files_parsing_code='\n'.join(output_files_parsing_code_lines),