diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 719cdfd7b0..0d4003ed28 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -5,6 +5,7 @@ from google.protobuf import struct_pb2 as _struct from flytekit.models import common as _common +from flytekit.models.annotation import TypeAnnotation as _annotation_model from flytekit.models.core import types as _core_types @@ -108,6 +109,7 @@ def __init__( blob=None, enum_type=None, metadata=None, + annotation=None, ): """ This is a oneof message, only one of the kwargs may be set, representing one of the Flyte types. @@ -120,6 +122,8 @@ def __init__( :param flytekit.models.core.types.BlobType blob: For blob objects, this describes the type. :param flytekit.models.core.types.EnumType enum_type: For enum objects, describes an enum :param dict[Text, T] metadata: Additional data describing the type + :param flytekit.models.annotation.FlyteAnnotation annotation: Additional data + describing the type _intended to be saturated by the client_ """ self._simple = simple self._schema = schema @@ -128,6 +132,7 @@ def __init__( self._blob = blob self._enum_type = enum_type self._metadata = metadata + self._annotation = annotation @property def simple(self) -> SimpleType: @@ -166,18 +171,31 @@ def metadata(self): """ return self._metadata + @property + def annotation(self) -> _annotation_model: + """ + :rtype: flytekit.models.annotation.FlyteAnnotation + """ + return self._annotation + @metadata.setter def metadata(self, value): self._metadata = value + @annotation.setter + def annotation(self, value): + self.annotation = value + def to_flyte_idl(self): """ :rtype: flyteidl.core.types_pb2.LiteralType """ + if self.metadata is not None: metadata = _json_format.Parse(_json.dumps(self.metadata), _struct.Struct()) else: metadata = None + t = _types_pb2.LiteralType( simple=self.simple if self.simple is not None else None, schema=self.schema.to_flyte_idl() if self.schema is not None else None, @@ -186,6 +204,7 @@ def to_flyte_idl(self): blob=self.blob.to_flyte_idl() if self.blob is not None else None, enum_type=self.enum_type.to_flyte_idl() if self.enum_type else None, metadata=metadata, + annotation=self.annotation.to_flyte_idl() if self.annotation else None, ) return t @@ -209,6 +228,7 @@ def from_flyte_idl(cls, proto): blob=_core_types.BlobType.from_flyte_idl(proto.blob) if proto.HasField("blob") else None, enum_type=_core_types.EnumType.from_flyte_idl(proto.enum_type) if proto.HasField("enum_type") else None, metadata=_json_format.MessageToDict(proto.metadata) or None, + annotation=_annotation_model.from_flyte_idl(proto.annotation) if proto.HasField("annotation") else None, )