Refactor extract_types in Dict Transformer

flytekit/core/type_engine.py

@@ -2019,25 +2019,41 @@ def __init__(self):
 
     @staticmethod
     def extract_types(t: Optional[Type[dict]]) -> typing.Tuple:
+        if t is None:
+            return None, None
+
+        # Get the origin and type arguments.
         _origin = get_origin(t)
         _args = get_args(t)
-        if _origin is not None:
-            if _origin is Annotated and _args:
-                # _args holds the type arguments to the dictionary, in other words:
-                # >>> get_args(Annotated[dict[int, str], FlyteAnnotation("abc")])
-                # (dict[int, str], <flytekit.core.annotation.FlyteAnnotation object at 0x107f6ff80>)
-                for x in _args[1:]:
-                    if isinstance(x, FlyteAnnotation):
-                        raise ValueError(
-                            f"Flytekit does not currently have support for FlyteAnnotations applied to dicts. {t} cannot be parsed."
-                        )
-            if _origin is dict and _args is not None:
+
+        # If not annotated or dict, return None, None.
+        if _origin is None:
+            return None, None
+
+        # If this is something like Annotated[dict[int, str], FlyteAnnotation("abc")],
+        # we need to check if there's a FlyteAnnotation in the metadata.
+        if _origin is Annotated:
+            if not _args:
+                return None, None
+
+            first_arg = _args[0]
+            # Check the rest of the metadata (after the dict type itself).
+            for x in _args[1:]:
+                if isinstance(x, FlyteAnnotation):
+                    raise ValueError(
+                        f"Flytekit does not currently have support for FlyteAnnotations applied to dicts. {t} cannot be parsed."
+                    )
+            # Recursively process the first argument if it's Annotated (or dict).
+            return DictTransformer.extract_types(first_arg)
+
+        # If the origin is dict, return the type arguments if they exist.
+        if _origin is dict:
+            # _args can be ().
+            if _args is not None:
                 return _args  # type: ignore
-            elif _origin is Annotated:
-                return DictTransformer.extract_types(_args[0])
-            else:
-                raise ValueError(f"Trying to extract dictionary type information from a non-dict type {t}")
-        return None, None
+
+        # Otherwise, we do not support this type in extract_types.
+        raise ValueError(f"Trying to extract dictionary type information from a non-dict type {t}")
 
     @staticmethod
     async def dict_to_generic_literal(

tests/flytekit/unit/core/test_type_engine.py

@@ -2864,21 +2864,22 @@ def test_get_underlying_type(t, expected):
 
 
 @pytest.mark.parametrize(
-    "t,expected",
+    "t,expected,allow_pickle",
     [
-        (None, (None, None)),
-        (typing.Dict, ()),
-        (typing.Dict[str, str], (str, str)),
+        (None, (None, None), False),
+        (typing.Dict, (), False),
+        (typing.Dict[str, str], (str, str), False),
         (
-                Annotated[typing.Dict[str, str], kwtypes(allow_pickle=True)],
-                (str, str),
+            Annotated[typing.Dict[str, str], kwtypes(allow_pickle=True)],
+            (str, str),
+            True,
         ),
-        (typing.Dict[Annotated[str, "a-tag"], int], (Annotated[str, "a-tag"], int)),
+        (typing.Dict[Annotated[str, "a-tag"], int], (Annotated[str, "a-tag"], int), False),
     ],
 )
-def test_dict_get(t, expected):
+def test_dict_get(t, expected, allow_pickle):
     assert DictTransformer.extract_types(t) == expected
-
+    assert DictTransformer.is_pickle(t) == allow_pickle
 
 def test_DataclassTransformer_get_literal_type():
     @dataclass
@@ -3777,3 +3778,35 @@ class RegularDC:
     assert TypeEngine.get_transformer(RegularDC) == TypeEngine._DATACLASS_TRANSFORMER
 
     del TypeEngine._REGISTRY[ParentDC]
+
+
+@pytest.mark.asyncio
+async def test_dict_transformer_annotated_type():
+    ctx = FlyteContext.current_context()
+
+    # Test case 1: Regular Dict type
+    regular_dict = {"a": 1, "b": 2}
+    regular_dict_type = Dict[str, int]
+    expected_type = TypeEngine.to_literal_type(regular_dict_type)
+
+    # This should work fine
+    literal1 = await TypeEngine.async_to_literal(ctx, regular_dict, regular_dict_type, expected_type)
+    assert literal1.map.literals["a"].scalar.primitive.integer == 1
+    assert literal1.map.literals["b"].scalar.primitive.integer == 2
+
+    # Test case 2: Annotated Dict type
+    annotated_dict = {"x": 10, "y": 20}
+    annotated_dict_type = Annotated[Dict[str, int], "some_metadata"]
+    expected_type = TypeEngine.to_literal_type(annotated_dict_type)
+
+    literal2 = await TypeEngine.async_to_literal(ctx, annotated_dict, annotated_dict_type, expected_type)
+    assert literal2.map.literals["x"].scalar.primitive.integer == 10
+    assert literal2.map.literals["y"].scalar.primitive.integer == 20
+
+    # Test case 3: Nested Annotated Dict type
+    nested_dict = {"outer": {"inner": 42}}
+    nested_dict_type = Dict[str, Annotated[Dict[str, int], "inner_metadata"]]
+    expected_type = TypeEngine.to_literal_type(nested_dict_type)
+
+    literal3 = await TypeEngine.async_to_literal(ctx, nested_dict, nested_dict_type, expected_type)
+    assert literal3.map.literals["outer"].map.literals["inner"].scalar.primitive.integer == 42