fix(sdk.v2): fix a small bug in io_types.is_artifact_annotation() (#…

…5699)

* fix is_artifact_annotation

* move tests

sdk/python/kfp/dsl/io_types.py

@@ -432,6 +432,9 @@ def is_artifact_annotation(typ) -> bool:
   if not hasattr(typ, '__args__') or len(typ.__args__) != 2:
     return False
 
+  if typ.__args__[1] not in [InputAnnotation, OutputAnnotation]:
+    return False
+
   return True
 
 def is_input_artifact(typ) -> bool:

sdk/python/tests/dsl/io_types_test.py → sdk/python/kfp/dsl/io_types_test.py

@@ -16,6 +16,7 @@
 import unittest
 import json
 import os
+from typing import List, Optional, Union
 
 from kfp.dsl import io_types
 from kfp.dsl.io_types import Input, InputAnnotation, Output, Model, OutputAnnotation
@@ -31,34 +32,40 @@ def test_complex_metrics(self):
     metrics.log_confusion_matrix_row('dog', [2, 6, 0])
     metrics.log_confusion_matrix_cell('cat', 'dog', 3)
     metrics.log_confusion_matrix_cell('horses', 'horses', 3)
-    metrics.metadata["test"] = 1.0
-    with open(os.path.join(os.path.dirname(__file__),
-        'test_data', 'expected_io_types_classification_metrics.json')) as json_file:
+    metrics.metadata['test'] = 1.0
+    with open(
+        os.path.join(
+            os.path.dirname(__file__), 'test_data',
+            'expected_io_types_classification_metrics.json')) as json_file:
       expected_json = json.load(json_file)
       self.assertEqual(expected_json, metrics.metadata)
 
   def test_complex_metrics_bulk_loading(self):
     metrics = io_types.ClassificationMetrics()
-    metrics.log_roc_curve(fpr=[85.1, 85.1, 85.1],
-                          tpr=[52.6, 52.6, 52.6],
-                          threshold=[53.6, 53.6, 53.6])
+    metrics.log_roc_curve(
+        fpr=[85.1, 85.1, 85.1],
+        tpr=[52.6, 52.6, 52.6],
+        threshold=[53.6, 53.6, 53.6])
     metrics.log_confusion_matrix(['dog', 'cat', 'horses'],
                                  [[2, 6, 0], [3, 5, 6], [5, 7, 8]])
-    with open(os.path.join(os.path.dirname(__file__),
-        'test_data', 'expected_io_types_bulk_load_classification_metrics.json')) as json_file:
+    with open(
+        os.path.join(
+            os.path.dirname(__file__), 'test_data',
+            'expected_io_types_bulk_load_classification_metrics.json')
+    ) as json_file:
       expected_json = json.load(json_file)
       self.assertEqual(expected_json, metrics.metadata)
 
-
-class InputOutputArtifacts(unittest.TestCase):
-
   def test_is_artifact_annotation(self):
     self.assertTrue(io_types.is_artifact_annotation(Input[Model]))
     self.assertTrue(io_types.is_artifact_annotation(Output[Model]))
+    self.assertTrue(io_types.is_artifact_annotation(Output['MyArtifact']))
 
     self.assertFalse(io_types.is_artifact_annotation(Model))
     self.assertFalse(io_types.is_artifact_annotation(int))
     self.assertFalse(io_types.is_artifact_annotation('Dataset'))
+    self.assertFalse(io_types.is_artifact_annotation(List[str]))
+    self.assertFalse(io_types.is_artifact_annotation(Optional[str]))
 
   def test_is_input_artifact(self):
     self.assertTrue(io_types.is_input_artifact(Input[Model]))
@@ -83,18 +90,18 @@ def test_get_io_artifact_class(self):
     self.assertEqual(io_types.get_io_artifact_class(str), None)
 
   def test_get_io_artifact_annotation(self):
-    self.assertEqual(io_types.get_io_artifact_annotation(Output[Model]),
-                     OutputAnnotation)
-    self.assertEqual(io_types.get_io_artifact_annotation(Input[Model]),
-                     InputAnnotation)
-    self.assertEqual(io_types.get_io_artifact_annotation(Input),
-                     InputAnnotation)
-    self.assertEqual(io_types.get_io_artifact_annotation(Output),
-                     OutputAnnotation)
+    self.assertEqual(
+        io_types.get_io_artifact_annotation(Output[Model]), OutputAnnotation)
+    self.assertEqual(
+        io_types.get_io_artifact_annotation(Input[Model]), InputAnnotation)
+    self.assertEqual(
+        io_types.get_io_artifact_annotation(Input), InputAnnotation)
+    self.assertEqual(
+        io_types.get_io_artifact_annotation(Output), OutputAnnotation)
 
     self.assertEqual(io_types.get_io_artifact_annotation(Model), None)
     self.assertEqual(io_types.get_io_artifact_annotation(str), None)
 
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

sdk/python/tests/dsl/metric_utils_test.py → sdk/python/kfp/dsl/metrics_utils_test.py

@@ -26,26 +26,30 @@
 
 from google.protobuf import json_format
 
+
 class MetricsUtilsTest(unittest.TestCase):
 
   def test_confusion_matrix(self):
     conf_matrix = metrics_utils.ConfusionMatrix()
     conf_matrix.set_categories(['dog', 'cat', 'horses'])
     conf_matrix.log_row('dog', [2, 6, 0])
     conf_matrix.log_cell('cat', 'dog', 3)
-    with open(os.path.join(os.path.dirname(__file__),
-        'test_data', 'expected_confusion_matrix.json')) as json_file:
+    with open(
+        os.path.join(
+            os.path.dirname(__file__), 'test_data',
+            'expected_confusion_matrix.json')) as json_file:
       expected_json = json.load(json_file)
       self.assertEqual(expected_json, conf_matrix.get_metrics())
 
   def test_bulkload_confusion_matrix(self):
     conf_matrix = metrics_utils.ConfusionMatrix()
-    conf_matrix.load_matrix(['dog', 'cat', 'horses'], [
-      [2, 6, 0], [3, 5,6], [5,7,8]])
+    conf_matrix.load_matrix(['dog', 'cat', 'horses'],
+                            [[2, 6, 0], [3, 5, 6], [5, 7, 8]])
 
-    with open(os.path.join(os.path.dirname(__file__),
-        'test_data',
-        'expected_bulk_loaded_confusion_matrix.json')) as json_file:
+    with open(
+        os.path.join(
+            os.path.dirname(__file__), 'test_data',
+            'expected_bulk_loaded_confusion_matrix.json')) as json_file:
       expected_json = json.load(json_file)
       self.assertEqual(expected_json, conf_matrix.get_metrics())
 
@@ -55,11 +59,12 @@ def test_confidence_metrics(self):
     confid_metrics.recall = 24.5
     confid_metrics.falsePositiveRate = 98.4
     expected_dict = {
-      'confidenceThreshold': 24.3,
-      'recall': 24.5,
-      'falsePositiveRate': 98.4
+        'confidenceThreshold': 24.3,
+        'recall': 24.5,
+        'falsePositiveRate': 98.4
     }
     self.assertEqual(expected_dict, confid_metrics.get_metrics())
 
+
 if __name__ == '__main__':
   unittest.main()