diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 5994390c8d..e01cb9a343 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -173,8 +173,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp return self._to_literal_transformer(python_val) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: - if get_origin(expected_python_type) is Annotated: - expected_python_type = get_args(expected_python_type)[0] + expected_python_type = get_underlying_type(expected_python_type) if expected_python_type != self._type: raise TypeTransformerFailedError( @@ -311,7 +310,7 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: Extracts the Literal type definition for a Dataclass and returns a type Struct. If possible also extracts the JSONSchema for the dataclass. """ - if get_origin(t) is Annotated: + if is_annotated(t): raise ValueError( "Flytekit does not currently have support for FlyteAnnotations applied to Dataclass." f"Type {t} cannot be parsed." @@ -368,7 +367,7 @@ def _get_origin_type_in_annotation(self, python_type: Type[T]) -> Type[T]: self._get_origin_type_in_annotation(get_args(python_type)[0]), self._get_origin_type_in_annotation(get_args(python_type)[1]), ] - elif get_origin(python_type) is Annotated: + elif is_annotated(python_type): return get_args(python_type)[0] elif dataclasses.is_dataclass(python_type): for field in dataclasses.fields(copy.deepcopy(python_type)): @@ -737,7 +736,7 @@ def get_transformer(cls, python_type: Type) -> TypeTransformer[T]: """ cls.lazy_import_transformers() # Step 1 - if get_origin(python_type) is Annotated: + if is_annotated(python_type): args = get_args(python_type) for annotation in args: if isinstance(annotation, TypeTransformer): @@ -752,7 +751,7 @@ def get_transformer(cls, python_type: Type) -> TypeTransformer[T]: if hasattr(python_type, "__origin__"): # Handling of annotated generics, eg: # Annotated[typing.List[int], 'foo'] - if get_origin(python_type) is Annotated: + if is_annotated(python_type): return cls.get_transformer(get_args(python_type)[0]) if python_type.__origin__ in cls._REGISTRY: @@ -823,7 +822,7 @@ def to_literal_type(cls, python_type: Type) -> LiteralType: transformer = cls.get_transformer(python_type) res = transformer.get_literal_type(python_type) data = None - if get_origin(python_type) is Annotated: + if is_annotated(python_type): for x in get_args(python_type)[1:]: if not isinstance(x, FlyteAnnotation): continue @@ -851,9 +850,9 @@ def to_literal(cls, ctx: FlyteContext, python_val: typing.Any, python_type: Type # In case the value is an annotated type we inspect the annotations and look for hash-related annotations. hash = None - if get_origin(python_type) is Annotated: + if is_annotated(python_type): # We are now dealing with one of two cases: - # 1. The annotated type is a `HashMethod`, which indicates that we should we should produce the hash using + # 1. The annotated type is a `HashMethod`, which indicates that we should produce the hash using # the method indicated in the annotation. # 2. The annotated type is being used for a different purpose other than calculating hash values, in which case # we should just continue. @@ -880,7 +879,7 @@ def to_python_value(cls, ctx: FlyteContext, lv: Literal, expected_python_type: T @classmethod def to_html(cls, ctx: FlyteContext, python_val: typing.Any, expected_python_type: Type[typing.Any]) -> str: transformer = cls.get_transformer(expected_python_type) - if get_origin(expected_python_type) is Annotated: + if is_annotated(expected_python_type): expected_python_type, *annotate_args = get_args(expected_python_type) from flytekit.deck.renderer import Renderable @@ -1004,7 +1003,7 @@ def get_sub_type(t: Type[T]) -> Type[T]: if hasattr(t, "__origin__"): # Handle annotation on list generic, eg: # Annotated[typing.List[int], 'foo'] - if get_origin(t) is Annotated: + if is_annotated(t): return ListTransformer.get_sub_type(get_args(t)[0]) if getattr(t, "__origin__") is list and hasattr(t, "__args__"): @@ -1030,7 +1029,7 @@ def is_batchable(t: Type): """ from flytekit.types.pickle import FlytePickle - if get_origin(t) is Annotated: + if is_annotated(t): return ListTransformer.is_batchable(get_args(t)[0]) if get_origin(t) is list: subtype = get_args(t)[0] @@ -1047,7 +1046,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp batch_size = len(python_val) # default batch size # parse annotated to get the number of items saved in a pickle file. - if get_origin(python_type) is Annotated: + if is_annotated(python_type): for annotation in get_args(python_type)[1:]: if isinstance(annotation, BatchSize): batch_size = annotation.val @@ -1191,8 +1190,7 @@ def get_sub_type_in_optional(t: Type[T]) -> Type[T]: return get_args(t)[0] def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: - if get_origin(t) is Annotated: - t = get_args(t)[0] + t = get_underlying_type(t) try: trans: typing.List[typing.Tuple[TypeTransformer, typing.Any]] = [ @@ -1206,8 +1204,7 @@ def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: 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: - if get_origin(python_type) is Annotated: - python_type = get_args(python_type)[0] + python_type = get_underlying_type(python_type) found_res = False res = None @@ -1232,8 +1229,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp raise TypeTransformerFailedError(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 get_origin(expected_python_type) is Annotated: - expected_python_type = get_args(expected_python_type)[0] + expected_python_type = get_underlying_type(expected_python_type) union_tag = None union_type = None @@ -1468,7 +1464,7 @@ def __init__(self): super().__init__(name="DefaultEnumTransformer", t=enum.Enum) def get_literal_type(self, t: Type[T]) -> LiteralType: - if get_origin(t) is Annotated: + if is_annotated(t): raise ValueError( f"Flytekit does not currently have support \ for FlyteAnnotations applied to enums. {t} cannot be \ @@ -1782,3 +1778,14 @@ def get(self, attr: str, as_type: Optional[typing.Type] = None) -> typing.Any: _register_default_type_transformers() + + +def is_annotated(t: Type) -> bool: + return get_origin(t) is Annotated + + +def get_underlying_type(t: Type) -> Type: + """Return the underlying type for annotated types or the type itself""" + if is_annotated(t): + return get_args(t)[0] + return t diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index bb8feb3d9c..d78ec152d7 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -8,10 +8,9 @@ from dataclasses_json import config, dataclass_json from marshmallow import fields -from typing_extensions import Annotated, get_args, get_origin from flytekit.core.context_manager import FlyteContext, FlyteContextManager -from flytekit.core.type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError +from flytekit.core.type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError, get_underlying_type from flytekit.loggers import logger from flytekit.models.core.types import BlobType from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar @@ -337,8 +336,7 @@ def to_literal( raise TypeTransformerFailedError("None value cannot be converted to a file.") # Correctly handle `Annotated[FlyteFile, ...]` by extracting the origin type - if get_origin(python_type) is Annotated: - python_type = get_args(python_type)[0] + python_type = get_underlying_type(python_type) if not (python_type is os.PathLike or issubclass(python_type, FlyteFile)): raise ValueError(f"Incorrect type {python_type}, must be either a FlyteFile or os.PathLike") @@ -413,6 +411,9 @@ def to_python_value( if expected_python_type is os.PathLike: return FlyteFile(uri) + # Correctly handle `Annotated[FlyteFile, ...]` by extracting the origin type + expected_python_type = get_underlying_type(expected_python_type) + # The rest of the logic is only for FlyteFile types. if not issubclass(expected_python_type, FlyteFile): # type: ignore raise TypeError(f"Neither os.PathLike nor FlyteFile specified {expected_python_type}") diff --git a/tests/flytekit/unit/core/test_flyte_file.py b/tests/flytekit/unit/core/test_flyte_file.py index b7f0a1aeee..5dd05cdffd 100644 --- a/tests/flytekit/unit/core/test_flyte_file.py +++ b/tests/flytekit/unit/core/test_flyte_file.py @@ -439,13 +439,20 @@ def test_flyte_file_annotated_hashmethod(local_dummy_file): def calc_hash(ff: FlyteFile) -> str: return str(ff.path) + HashedFlyteFile = Annotated[FlyteFile, HashMethod(calc_hash)] + @task - def t1(path: str) -> Annotated[FlyteFile, HashMethod(calc_hash)]: - return FlyteFile(path) + def t1(path: str) -> HashedFlyteFile: + return HashedFlyteFile(path) + + @task + def t2(ff: HashedFlyteFile) -> None: + print(ff.path) @workflow def wf(path: str) -> None: - t1(path=path) + ff = t1(path=path) + t2(ff=ff) wf(path=local_dummy_file) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 6d1b6829d5..2e52fdcf9d 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -41,6 +41,8 @@ UnionTransformer, convert_json_schema_to_python_class, dataclass_from_dict, + get_underlying_type, + is_annotated, ) from flytekit.exceptions import user as user_exceptions from flytekit.models import types as model_types @@ -1685,3 +1687,28 @@ def test_batch_pickle_list(python_val, python_type, expected_list_length): # data = task0() # task0() -> Annotated[typing.List[FlytePickle], BatchSize(2)] # task1(data=data) # task1(data: typing.List[FlytePickle]) assert pv == python_val + + +@pytest.mark.parametrize( + "t,expected", + [ + (list, False), + (Annotated[int, "tag"], True), + (Annotated[typing.List[str], "a", "b"], True), + (Annotated[typing.Dict[int, str], FlyteAnnotation({"foo": "bar"})], True), + ], +) +def test_is_annotated(t, expected): + assert is_annotated(t) == expected + + +@pytest.mark.parametrize( + "t,expected", + [ + (typing.List, typing.List), + (Annotated[int, "tag"], int), + (Annotated[typing.List[str], "a", "b"], typing.List[str]), + ], +) +def test_get_underlying_type(t, expected): + assert get_underlying_type(t) == expected