fixes handling of union types

In the dataclasses conversion code type that is a member of a union
was not properly checked for if it was a member and so there would
always be an error.

For instance `FlyteFile.path` is `Union[str,Pathlike]` and so `str !=
Union[str,Pathlike]`.

This patch adds support for checking that a type is part of a union
and a satisfactory type.

Signed-off-by: Samuel Lotz <[email protected]>

Author: salotz

flytekit/core/type_engine.py

@@ -17,7 +17,7 @@
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from functools import lru_cache
-from types import GenericAlias
+from types import GenericAlias, UnionType
 from typing import Any, Dict, List, NamedTuple, Optional, Type, cast
 
 import msgpack
@@ -565,7 +565,14 @@ def assert_type(self, expected_type: Type[DataClassJsonMixin], v: T):
                         original_type = type(v)
                         if UnionTransformer.is_optional_type(expected_type):
                             expected_type = UnionTransformer.get_sub_type_in_optional(expected_type)
-                        if original_type != expected_type:
+
+                        if (
+                                UnionTransformer.is_union(expected_type) and
+                                UnionTransformer.in_union(original_type, expected_type)
+                        ):
+                            pass
+                        
+                        elif original_type != expected_type:
                             raise TypeTransformerFailedError(
                                 f"Type of Val '{original_type}' is not an instance of {expected_type}"
                             )
@@ -1836,6 +1843,16 @@ class UnionTransformer(AsyncTypeTransformer[T]):
     def __init__(self):
         super().__init__("Typed Union", typing.Union)
 
+    @staticmethod
+    def is_union(t: Type[Any] | UnionType) -> bool:
+
+        return _is_union_type(t)
+
+    @staticmethod
+    def in_union(t: Type[Any], union: types.UnionType) -> bool:
+
+        return t in typing.get_args(union)
+        
     @staticmethod
     def is_optional_type(t: Type) -> bool:
         return _is_union_type(t) and type(None) in get_args(t)