From 189a5796ca1d0bfa03db76150a2f0b3d96bd534c Mon Sep 17 00:00:00 2001
From: connor-mccarthy <mccarthy.connor.james@gmail.com>
Date: Thu, 1 Sep 2022 15:32:06 -0600
Subject: [PATCH] use _from_executor_fields when available

---
 sdk/python/kfp/components/executor.py      |  6 ++++++
 sdk/python/kfp/components/executor_test.py | 17 ++++++++++++-----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/sdk/python/kfp/components/executor.py b/sdk/python/kfp/components/executor.py
index 332c7155485..99cbe908adb 100644
--- a/sdk/python/kfp/components/executor.py
+++ b/sdk/python/kfp/components/executor.py
@@ -326,6 +326,12 @@ def create_artifact_instance(
 
     artifact_cls = artifact_types._SCHEMA_TITLE_TO_TYPE.get(
         schema_title) or artifact_cls or artifact_types.Artifact
+    if hasattr(artifact_cls, '_from_executor_fields'):
+        return artifact_cls._from_executor_fields(
+            uri=runtime_artifact.get('uri', ''),
+            name=runtime_artifact.get('name', ''),
+            metadata=runtime_artifact.get('metadata', {}),
+        )
     return artifact_cls(
         uri=runtime_artifact.get('uri', ''),
         name=runtime_artifact.get('name', ''),
diff --git a/sdk/python/kfp/components/executor_test.py b/sdk/python/kfp/components/executor_test.py
index 6670dd757d5..0459934d169 100644
--- a/sdk/python/kfp/components/executor_test.py
+++ b/sdk/python/kfp/components/executor_test.py
@@ -826,10 +826,17 @@ class VertexDataset:
     schema_title = 'google.VertexDataset'
     schema_version = '0.0.0'
 
-    def __init__(self, name: str, uri: str, metadata: dict) -> None:
-        self.name = name
-        self.uri = uri
-        self.metadata = metadata
+    def __init__(self, display_name: str) -> None:
+        self.display_name = display_name
+        self.name = 'my_name'
+        self.uri = 'my_uri'
+        self.metadata = {}
+
+    @classmethod
+    def _from_executor_fields(cls, name: str, uri: str,
+                              metadata: dict) -> 'VertexDataset':
+        print('Used _from_executor_fields to construct')
+        return VertexDataset(display_name=name)
 
     @property
     def path(self) -> str:
@@ -960,7 +967,7 @@ def test_dict_to_artifact_nonkfp_artifact(self):
             },
             'uri': 'gs://some-bucket/input_artifact_one'
         }
-        # with artifact_cls
+        # with artifact_cls, using _from_executor_fields
         self.assertIsInstance(
             executor.create_artifact_instance(
                 runtime_artifact, artifact_cls=VertexDataset), VertexDataset)