Add support union type

.github/workflows/pythonbuild.yml

@@ -36,6 +36,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip==21.2.4 setuptools wheel
         make setup${{ matrix.spark-version-suffix }}
+        git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd ..
         pip freeze
     - name: Test with coverage
       run: |
@@ -93,6 +94,7 @@ jobs:
         cd plugins/${{ matrix.plugin-names }}
         pip install -e .
         pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit
+        git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd ..
         pip freeze
     - name: Test with coverage
       run: |
@@ -140,5 +142,6 @@ jobs:
         run: |
           python -m pip install --upgrade pip==21.2.4 setuptools wheel
           pip install -r doc-requirements.txt
+          git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd ..
       - name: Build the documentation
         run: make -C docs html

flytekit/core/type_engine.py

@@ -28,7 +28,7 @@
 from flytekit.models import types as _type_models
 from flytekit.models.core import types as _core_types
 from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar
-from flytekit.models.types import LiteralType, SimpleType
+from flytekit.models.types import LiteralType, SimpleType, UnionType
 
 T = typing.TypeVar("T")
 DEFINITIONS = "definitions"
@@ -444,7 +444,7 @@ def to_literal(cls, ctx: FlyteContext, python_val: typing.Any, python_type: Type
         """
         Converts a python value of a given type and expected ``LiteralType`` into a resolved ``Literal`` value.
         """
-        if python_val is None:
+        if python_val is None and expected.union_type is None:
             raise AssertionError(f"Python value cannot be None, expected {python_type}/{expected}")
         transformer = cls.get_transformer(python_type)
         if transformer.type_assertions_enabled:
@@ -596,6 +596,58 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[list]:
         raise ValueError(f"List transformer cannot reverse {literal_type}")
 
 
+class UnionTransformer(TypeTransformer[T]):
+    """
+    Transformer that handles a univariate typing.Union[T]
+    """
+
+    def __init__(self):
+        super().__init__("Typed Union", typing.Union)
+
+    @staticmethod
+    def get_sub_type(t: Type[T]) -> Type[T]:
+        """
+        Return the generic Type T of the Union
+        """
+        if hasattr(t, "__origin__") and t.__origin__ is typing.Union:  # type: ignore
+            if hasattr(t, "__args__"):
+                return t.__args__  # type: ignore
+        raise ValueError("Only generic univariate typing.Union[T] type is supported.")
+
+    def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]:
+        try:
+            sub_type = [TypeEngine.to_literal_type(v) if v else "" for v in self.get_sub_type(t)]
+            return _type_models.LiteralType(union_type=UnionType(sub_type))
+        except Exception as e:
+            raise ValueError(f"Type of Generic Union type is not supported, {e}")
+
+    def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal:
+        for t in python_type.__args__:
+            try:
+                return TypeEngine.to_literal(ctx, python_val, t, expected)
+            except Exception as e:
+                logger.debug(f"Failed to convert from {python_val} to {t}", e)
+        raise TypeError(f"Cannot convert from {python_val} to {python_type}")
+
+    def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]:
+        if lv is None:
+            return None
+        st = self.get_sub_type(expected_python_type)
+        for v in st:
+            try:
+                val = TypeEngine.to_python_value(ctx, lv, v)
+                if val:
+                    return val
+            except Exception as e:
+                logger.debug(f"Failed to convert from {lv} to {v}", e)
+        raise TypeError(f"Cannot convert from {lv} to {expected_python_type}")
+
+    def guess_python_type(self, literal_type: LiteralType) -> type:
+        if literal_type.union_type:
+            return typing.Union[tuple(TypeEngine.guess_python_type(v) for v in literal_type.union_type.values)]
+        raise ValueError(f"Union transformer cannot reverse {literal_type}")
+
+
 class DictTransformer(TypeTransformer[dict]):
     """
     Transformer that transforms a univariate dictionary Dict[str, T] to a Literal Map or
@@ -907,9 +959,11 @@ def _register_default_type_transformers():
             _type_models.LiteralType(simple=_type_models.SimpleType.NONE),
             lambda x: None,
             lambda x: None,
-        )
+        ),
+        [type(None)],
     )
     TypeEngine.register(ListTransformer())
+    TypeEngine.register(UnionTransformer())
     TypeEngine.register(DictTransformer())
     TypeEngine.register(TextIOTransformer())
     TypeEngine.register(BinaryIOTransformer())

flytekit/models/types.py

@@ -1,4 +1,5 @@
 import json as _json
+import typing
 
 from flyteidl.core import types_pb2 as _types_pb2
 from google.protobuf import json_format as _json_format
@@ -98,6 +99,28 @@ def from_flyte_idl(cls, proto):
         return cls(columns=[SchemaType.SchemaColumn.from_flyte_idl(c) for c in proto.columns])
 
 
+class UnionType(_common.FlyteIdlEntity):
+    """
+    Models _types_pb2.UnionType
+    """
+
+    def __init__(self, values: typing.List["LiteralType"]):
+        self._values = values
+
+    @property
+    def values(self) -> typing.List["LiteralType"]:
+        return self._values
+
+    def to_flyte_idl(self) -> _types_pb2.UnionType:
+        return _types_pb2.UnionType(
+            values=[val.to_flyte_idl() if val else None for val in self._values],
+        )
+
+    @classmethod
+    def from_flyte_idl(cls, proto: _types_pb2.UnionType):
+        return cls(values=[LiteralType.from_flyte_idl(v) for v in proto.values])
+
+
 class LiteralType(_common.FlyteIdlEntity):
     def __init__(
         self,
@@ -107,6 +130,7 @@ def __init__(
         map_value_type=None,
         blob=None,
         enum_type=None,
+        union_type=None,
         metadata=None,
     ):
         """
@@ -119,6 +143,7 @@ def __init__(
             string.
         :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 flytekit.models.core.types.UnionType union_type: For union objects, describes an python union type.
         :param dict[Text, T] metadata: Additional data describing the type
         """
         self._simple = simple
@@ -127,6 +152,7 @@ def __init__(
         self._map_value_type = map_value_type
         self._blob = blob
         self._enum_type = enum_type
+        self._union_type = union_type
         self._metadata = metadata
 
     @property
@@ -159,6 +185,10 @@ def blob(self) -> _core_types.BlobType:
     def enum_type(self) -> _core_types.EnumType:
         return self._enum_type
 
+    @property
+    def union_type(self) -> UnionType:
+        return self._union_type
+
     @property
     def metadata(self):
         """
@@ -185,6 +215,7 @@ def to_flyte_idl(self):
             map_value_type=self.map_value_type.to_flyte_idl() if self.map_value_type is not None else None,
             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,
+            union_type=self.union_type.to_flyte_idl() if self.union_type else None,
             metadata=metadata,
         )
         return t
@@ -208,6 +239,7 @@ def from_flyte_idl(cls, proto):
             map_value_type=map_value_type,
             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,
+            union_type=UnionType.from_flyte_idl(proto.union_type) if proto.HasField("union_type") else None,
             metadata=_json_format.MessageToDict(proto.metadata) or None,
         )
 

tests/flytekit/unit/core/test_type_engine.py

@@ -547,6 +547,39 @@ def test_enum_type():
         TypeEngine.to_literal_type(UnsupportedEnumValues)
 
 
+def test_union_type():
+    pt = typing.Union[str, int]
+    lt = TypeEngine.to_literal_type(pt)
+    assert lt.union_type.values == [LiteralType(simple=SimpleType.STRING), LiteralType(simple=SimpleType.INTEGER)]
+
+    ctx = FlyteContextManager.current_context()
+    lv = TypeEngine.to_literal(ctx, 3, pt, lt)
+    v = TypeEngine.to_python_value(ctx, lv, pt)
+    assert lv.scalar.primitive.integer == 3
+    assert v == 3
+
+    lv = TypeEngine.to_literal(ctx, "hello", pt, lt)
+    v = TypeEngine.to_python_value(ctx, lv, pt)
+    assert lv.scalar.primitive.string_value == "hello"
+    assert v == "hello"
+
+
+def test_optional_type():
+    pt = typing.Optional[int]
+    lt = TypeEngine.to_literal_type(pt)
+    assert lt.union_type.values == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.NONE)]
+
+    ctx = FlyteContextManager.current_context()
+    lv = TypeEngine.to_literal(ctx, 3, pt, lt)
+    v = TypeEngine.to_python_value(ctx, lv, pt)
+    assert lv.scalar.primitive.integer == 3
+    assert v == 3
+
+    lv = TypeEngine.to_literal(ctx, None, pt, lt)
+    v = TypeEngine.to_python_value(ctx, lv, pt)
+    assert v is None
+
+
 @pytest.mark.parametrize(
     "python_value,python_types,expected_literal_map",
     [

tests/flytekit/unit/core/test_type_hints.py

@@ -3,6 +3,7 @@
 import functools
 import os
 import random
+import tempfile
 import typing
 from collections import OrderedDict
 from dataclasses import dataclass
@@ -31,6 +32,7 @@
 from flytekit.models.interface import Parameter
 from flytekit.models.task import Resources as _resource_models
 from flytekit.models.types import LiteralType, SimpleType
+from flytekit.types.file import FlyteFile
 from flytekit.types.schema import FlyteSchema, SchemaOpenMode
 
 serialization_settings = context_manager.SerializationSettings(
@@ -1427,3 +1429,51 @@ def foo3(a: typing.Dict) -> typing.Dict:
 
     with pytest.raises(TypeError, match="Not a collection type simple: STRUCT\n but got a list \\[{'hello': 2}\\]"):
         foo3(a=[{"hello": 2}])
+
+
+def test_union_type():
+    ut = typing.Union[int, str, float, FlyteFile, FlyteSchema, typing.List[int], typing.Dict[str, int]]
+
+    @task
+    def t1(a: ut) -> ut:
+        return a
+
+    @workflow
+    def wf(a: ut) -> ut:
+        return t1(a=a)
+
+    assert wf(a=2) == 2
+    assert wf(a="2") == "2"
+    assert wf(a=2.0) == 2.0
+    file = tempfile.NamedTemporaryFile()
+    assert isinstance(wf(a=FlyteFile(file.name)), FlyteFile)
+    assert isinstance(wf(a=FlyteSchema()), FlyteSchema)
+    assert wf(a=[1, 2, 3]) == [1, 2, 3]
+    assert wf(a={"a": 1}) == {"a": 1}
+
+    @task
+    def t2(a: typing.Union[float, dict]) -> typing.Union[float, dict]:
+        return a
+
+    @workflow
+    def wf2(a: typing.Union[int, str]) -> typing.Union[int, str]:
+        return t2(a=a)
+
+    with pytest.raises(
+        TypeError,
+        match='Cannot convert from scalar {\n  primitive {\n    string_value: "2"\n  }\n}\n to typing.Union\\[float, dict\\]',
+    ):
+        assert wf2(a="2") == "2"
+
+
+def test_optional_type():
+    @task
+    def t1(a: typing.Optional[int]) -> typing.Optional[int]:
+        return a
+
+    @workflow
+    def wf(a: typing.Optional[int]) -> typing.Optional[int]:
+        return t1(a=a)
+
+    assert wf(a=2) == 2
+    assert wf(a=None) is None