diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index d02016a465..0187680965 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -476,16 +476,15 @@ def dict_to_literal_map( # to account for the type erasure that happens in the case of built-in collection containers, such as # `list` and `dict`. python_type = guessed_python_types.get(k, type(v)) - if (hasattr(python_type, "__origin__") and not isinstance(v, python_type.__origin__)) or ( - not hasattr(python_type, "__origin__") and not isinstance(v, python_type) - ): + try: + literal_map[k] = TypeEngine.to_literal( + ctx=ctx, + python_val=v, + python_type=python_type, + expected=TypeEngine.to_literal_type(python_type), + ) + except TypeError: raise user_exceptions.FlyteTypeException(type(v), python_type, received_value=v) - literal_map[k] = TypeEngine.to_literal( - ctx=ctx, - python_val=v, - python_type=python_type, - expected=TypeEngine.to_literal_type(python_type), - ) return LiteralMap(literal_map) @classmethod diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index cb9512dfee..74628bb415 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -28,6 +28,7 @@ from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar from flytekit.models.types import LiteralType, SimpleType from flytekit.types.directory.types import FlyteDirectory +from flytekit.types.file import JPEGImageFile from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer @@ -600,6 +601,24 @@ def test_enum_type(): } ), ), + ( + {"p1": "s3://tmp/file.jpeg"}, + {"p1": JPEGImageFile}, + LiteralMap( + literals={ + "p1": Literal( + scalar=Scalar( + blob=Blob( + metadata=BlobMetadata( + type=BlobType(format="jpeg", dimensionality=BlobType.BlobDimensionality.SINGLE) + ), + uri="s3://tmp/file.jpeg", + ) + ) + ) + } + ), + ), ], ) def test_dict_to_literal_map(python_value, python_types, expected_literal_map):