From 7b4393dbdfb011551e10c34e65f4247717a6b9f6 Mon Sep 17 00:00:00 2001
From: connor-mccarthy <mccarthy.connor.james@gmail.com>
Date: Tue, 27 Sep 2022 17:04:33 -0600
Subject: [PATCH] fix and add tests

---
 .../components/component_decorator_test.py    | 44 +++++++++++++++++++
 .../kfp/components/component_factory.py       | 18 +++-----
 2 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/sdk/python/kfp/components/component_decorator_test.py b/sdk/python/kfp/components/component_decorator_test.py
index 4778a2abc21..5ba32915e0d 100644
--- a/sdk/python/kfp/components/component_decorator_test.py
+++ b/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]')
diff --git a/sdk/python/kfp/components/component_factory.py b/sdk/python/kfp/components/component_factory.py
index 81f56abc287..752d6f67592 100644
--- a/sdk/python/kfp/components/component_factory.py
+++ b/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