diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 9d1ad92a54..8bc41f74d4 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -774,7 +774,15 @@ def get_transformer(cls, python_type: Type) -> TypeTransformer[T]: python_type = args[0] - if python_type in cls._REGISTRY: + # this makes sure that if it's a list/dict of annotated types, we hit the unwrapping code in step 2 + # see test_list_of_annotated in test_structured_dataset.py + if ( + (not hasattr(python_type, "__origin__")) + or ( + hasattr(python_type, "__origin__") + and (python_type.__origin__ is not list and python_type.__origin__ is not dict) + ) + ) and python_type in cls._REGISTRY: return cls._REGISTRY[python_type] # Step 2 diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py index eaba8b6343..5124193b27 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -493,3 +493,17 @@ def to_html(self, input: str) -> str: with pytest.raises(NotImplementedError, match="Could not find a renderer for in"): StructuredDatasetTransformerEngine().to_html(FlyteContextManager.current_context(), 3, int) + + +def test_list_of_annotated(): + WineDataset = Annotated[ + StructuredDataset, + kwtypes( + alcohol=float, + malic_acid=float, + ), + ] + + @task + def no_op(data: WineDataset) -> typing.List[WineDataset]: + return [data]