fix and add tests

sdk/python/kfp/components/component_decorator_test.py

@@ -14,6 +14,7 @@
 
 import os
 import tempfile
+from typing import Dict, List, NamedTuple
 import unittest
 
 from kfp.components import python_component
@@ -106,3 +107,46 @@ def comp(text: str) -> str:
         component_spec = structures.ComponentSpec.load_from_component_yaml(
             yaml_text)
         self.assertEqual(component_spec, comp.component_spec)
+
+    def test_output_named_tuple_with_dict(self):
+
+        @component
+        def comp(
+                text: str) -> NamedTuple('outputs', [('data', Dict[str, str])]):
+            outputs = NamedTuple('outputs', [('data', Dict[str, str])])
+            return outputs(data={text: text})
+
+        # TODO: ideally should be the canonical type string, rather than the specific annotation as string, but both work
+        self.assertEqual(comp.component_spec.outputs['data'].type,
+                         'typing.Dict[str, str]')
+
+    def test_output_dict(self):
+
+        @component
+        def comp(text: str) -> Dict[str, str]:
+            return {text: text}
+
+        # TODO: ideally should be the canonical type string, rather than the specific annotation as string, but both work
+        self.assertEqual(comp.component_spec.outputs['Output'].type,
+                         'typing.Dict[str, str]')
+
+    def test_output_named_tuple_with_list(self):
+
+        @component
+        def comp(text: str) -> NamedTuple('outputs', [('data', List[str])]):
+            outputs = NamedTuple('outputs', [('data', List[str])])
+            return outputs(data={text: text})
+
+        # TODO: ideally should be the canonical type string, rather than the specific annotation as string, but both work
+        self.assertEqual(comp.component_spec.outputs['data'].type,
+                         'typing.List[str]')
+
+    def test_output_list(self):
+
+        @component
+        def comp(text: str) -> List[str]:
+            return {text: text}
+
+        # TODO: ideally should be the canonical type string, rather than the specific annotation as string, but both work
+        self.assertEqual(comp.component_spec.outputs['Output'].type,
+                         'typing.List[str]')

sdk/python/kfp/components/component_factory.py

@@ -242,21 +242,17 @@ def extract_component_interface(
                                         None) or getattr(
                                             return_ann, '_field_types', None)
             for field_name in return_ann._fields:
-                type_struct = None
-                if field_annotations:
-                    type_struct = type_utils._annotation_to_type_struct(
-                        field_annotations.get(field_name, None))
-
                 output_name = _maybe_make_unique(field_name, output_names)
                 output_names.add(output_name)
-                if type_struct.lower() in type_utils._PARAMETER_TYPES_MAPPING:
-
-                    output_spec = structures.OutputSpec(type=type_struct)
-                else:
+                annotation = field_annotations.get(field_name)
+                if type_annotations.is_artifact_class(annotation):
                     output_spec = structures.OutputSpec(
                         type=type_utils.create_bundled_artifact_type(
-                            type_struct,
-                            field_annotations.get(field_name).schema_version))
+                            annotation.schema_title, annotation.schema_version))
+                else:
+                    type_struct = type_utils._annotation_to_type_struct(
+                        annotation)
+                    output_spec = structures.OutputSpec(type=type_struct)
                 outputs[output_name] = output_spec
         # Deprecated dict-based way of declaring multiple outputs. Was only used by
         # the @component decorator