Skip to content
This repository has been archived by the owner on Jul 19, 2024. It is now read-only.

Commit

Permalink
Fix list of annotated structured dataset (flyteorg#1817)
Browse files Browse the repository at this point in the history
Signed-off-by: Yee Hing Tong <[email protected]>
Signed-off-by: Future Outlier <[email protected]>
  • Loading branch information
wild-endeavor authored and Future Outlier committed Oct 3, 2023
1 parent 664ca67 commit c5a7158
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 1 deletion.
10 changes: 9 additions & 1 deletion flytekit/core/type_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 14 additions & 0 deletions tests/flytekit/unit/core/test_structured_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -493,3 +493,17 @@ def to_html(self, input: str) -> str:

with pytest.raises(NotImplementedError, match="Could not find a renderer for <class 'int'> 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]

0 comments on commit c5a7158

Please sign in to comment.