Fix list of annotated structured dataset (flyteorg#1817)

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

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

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 <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]