feat: addition of annotation att to

Signed-off-by: Kenny Workman <[email protected]>

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,
         )