def testArtifactTypeRoundTrip(self): mlmd_artifact_type = standard_artifacts.Examples._get_artifact_type() self.assertIs( standard_artifacts.Examples, artifact_utils.get_artifact_type_class(mlmd_artifact_type)) mlmd_artifact_type = _MyArtifact._get_artifact_type() self.assertIs( _MyArtifact, artifact_utils.get_artifact_type_class(mlmd_artifact_type))
def testArtifactTypeRoundTrip(self): mlmd_artifact_type = standard_artifacts.Examples._get_artifact_type() self.assertIs( standard_artifacts.Examples, artifact_utils.get_artifact_type_class(mlmd_artifact_type)) mlmd_artifact_type = _MyArtifact._get_artifact_type() # Test that the ID is ignored for type comparison purposes during # deserialization. mlmd_artifact_type.id = 123 self.assertIs( _MyArtifact, artifact_utils.get_artifact_type_class(mlmd_artifact_type))
def testArtifactTypeRoundTripUnknownArtifactClass(self, mock_warning): mlmd_artifact_type = copy.deepcopy( standard_artifacts.Examples._get_artifact_type()) self.assertIs( standard_artifacts.Examples, artifact_utils.get_artifact_type_class(mlmd_artifact_type)) mlmd_artifact_type.name = 'UnknownTypeName' reconstructed_class = artifact_utils.get_artifact_type_class( mlmd_artifact_type) mock_warning.assert_called_once() self.assertIsNot(standard_artifacts.Examples, reconstructed_class) self.assertTrue(issubclass(reconstructed_class, artifact.Artifact)) self.assertEqual('UnknownTypeName', reconstructed_class.TYPE_NAME) self.assertEqual(mlmd_artifact_type, reconstructed_class._get_artifact_type())
def from_json_dict(cls, dict_data: Dict[Text, Any]) -> Any: artifact_type = metadata_store_pb2.ArtifactType() json_format.Parse(json.dumps(dict_data['type']), artifact_type) type_cls = artifact_utils.get_artifact_type_class(artifact_type) artifacts = list( Artifact.from_json_dict(a) for a in dict_data['artifacts']) producer_component_id = dict_data.get('producer_component_id', None) output_key = dict_data.get('output_key', None) return Channel(type=type_cls, artifacts=artifacts, producer_component_id=producer_component_id, output_key=output_key)