Fix dict annotations (#3123)

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

Author: wild-endeavor

flytekit/core/promise.py

@@ -947,7 +947,7 @@ async def binding_data_from_python_std(
             if transformer_override and hasattr(transformer_override, "extract_types_or_metadata"):
                 _, v_type = transformer_override.extract_types_or_metadata(t_value_type)  # type: ignore
             else:
-                _, v_type = DictTransformer.extract_types_or_metadata(t_value_type)  # type: ignore
+                _, v_type = DictTransformer.extract_types(cast(typing.Type[dict], t_value_type))
             m = _literals_models.BindingDataMap(
                 bindings={
                     k: await binding_data_from_python_std(

flytekit/core/type_engine.py

@@ -878,7 +878,7 @@ def _fix_val_int(self, t: typing.Type, val: typing.Any) -> typing.Any:
             return list(map(lambda x: self._fix_val_int(ListTransformer.get_sub_type(t), x), val))
 
         if isinstance(val, dict):
-            ktype, vtype = DictTransformer.extract_types_or_metadata(t)
+            ktype, vtype = DictTransformer.extract_types(t)
             # Handle nested Dict. e.g. {1: {2: 3}, 4: {5: 6}})
             return {
                 self._fix_val_int(cast(type, ktype), k): self._fix_val_int(cast(type, vtype), v) for k, v in val.items()
@@ -2018,7 +2018,7 @@ def __init__(self):
         super().__init__("Typed Dict", dict)
 
     @staticmethod
-    def extract_types_or_metadata(t: Optional[Type[dict]]) -> typing.Tuple:
+    def extract_types(t: Optional[Type[dict]]) -> typing.Tuple:
         _origin = get_origin(t)
         _args = get_args(t)
         if _origin is not None:
@@ -2031,8 +2031,12 @@ def extract_types_or_metadata(t: Optional[Type[dict]]) -> typing.Tuple:
                         raise ValueError(
                             f"Flytekit does not currently have support for FlyteAnnotations applied to dicts. {t} cannot be parsed."
                         )
-            if _origin in [dict, Annotated] and _args is not None:
+            if _origin is dict and _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
 
     @staticmethod
@@ -2099,31 +2103,24 @@ async def dict_to_binary_literal(
             raise TypeTransformerFailedError(f"Cannot convert `{v}` to Flyte Literal.\n" f"Error Message: {e}")
 
     @staticmethod
-    def is_pickle(python_type: Type[dict]) -> typing.Tuple[bool, Type]:
-        base_type, *metadata = DictTransformer.extract_types_or_metadata(python_type)
+    def is_pickle(python_type: Type[dict]) -> bool:
+        _origin = get_origin(python_type)
+        metadata: typing.Tuple = ()
+        if _origin is Annotated:
+            metadata = get_args(python_type)[1:]
 
         for each_metadata in metadata:
             if isinstance(each_metadata, OrderedDict):
                 allow_pickle = each_metadata.get("allow_pickle", False)
-                return allow_pickle, base_type
-
-        return False, base_type
+                return allow_pickle
 
-    @staticmethod
-    def dict_types(python_type: Type) -> typing.Tuple[typing.Any, ...]:
-        if get_origin(python_type) is Annotated:
-            base_type, *_ = DictTransformer.extract_types_or_metadata(python_type)
-            tp = get_args(base_type)
-        else:
-            tp = DictTransformer.extract_types_or_metadata(python_type)
-
-        return tp
+        return False
 
     def get_literal_type(self, t: Type[dict]) -> LiteralType:
         """
         Transforms a native python dictionary to a flyte-specific ``LiteralType``
         """
-        tp = self.dict_types(t)
+        tp = DictTransformer.extract_types(t)
 
         if tp:
             if tp[0] == str:
@@ -2144,10 +2141,9 @@ async def async_to_literal(
             raise TypeTransformerFailedError("Expected a dict")
 
         allow_pickle = False
-        base_type = None
 
         if get_origin(python_type) is Annotated:
-            allow_pickle, base_type = DictTransformer.is_pickle(python_type)
+            allow_pickle = DictTransformer.is_pickle(python_type)
 
         if expected and expected.simple and expected.simple == SimpleType.STRUCT:
             if str2bool(os.getenv(FLYTE_USE_OLD_DC_FORMAT)):
@@ -2160,11 +2156,7 @@ async def async_to_literal(
                 raise ValueError("Flyte MapType expects all keys to be strings")
             # TODO: log a warning for Annotated objects that contain HashMethod
 
-            if base_type:
-                _, v_type = get_args(base_type)
-            else:
-                _, v_type = self.extract_types_or_metadata(python_type)
-
+            _, v_type = self.extract_types(python_type)
             lit_map[k] = TypeEngine.async_to_literal(ctx, v, cast(type, v_type), expected.map_value_type)
         vals = await _run_coros_in_chunks([c for c in lit_map.values()], batch_size=_TYPE_ENGINE_COROS_BATCH_SIZE)
         for idx, k in zip(range(len(vals)), lit_map.keys()):
@@ -2177,9 +2169,9 @@ async def async_to_python_value(self, ctx: FlyteContext, lv: Literal, expected_p
             return self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
 
         if lv and lv.map and lv.map.literals is not None:
-            tp = self.dict_types(expected_python_type)
+            tp = DictTransformer.extract_types(expected_python_type)
 
-            if tp is None or tp[0] is None:
+            if tp is None or len(tp) == 0 or tp[0] is None:
                 raise TypeError(
                     "TypeMismatch: Cannot convert to python dictionary from Flyte Literal Dictionary as the given "
                     "dictionary does not have sub-type hints or they do not match with the originating dictionary "

tests/flytekit/unit/core/test_annotated_bindings.py

@@ -1,12 +1,16 @@
 import asyncio
+from flytekit.core.artifact import Artifact
 from dataclasses import dataclass
-from typing import List, Optional, TypeVar, Type, Tuple
+from typing import List, Optional, TypeVar, Type, Tuple, Union
 
+from pydantic import BaseModel
 from typing_extensions import Annotated
 
-from flytekit import task, workflow
+from flytekit import task
+from flytekit import workflow
 from flytekit.configuration import SerializationSettings, ImageConfig, Image
-from flytekit.core.context_manager import FlyteContextManager, ExecutionState, FlyteContext
+from flytekit.core.context_manager import FlyteContext
+from flytekit.core.context_manager import FlyteContextManager, ExecutionState
 from flytekit.core.dynamic_workflow_task import dynamic
 from flytekit.core.type_engine import SimpleTransformer
 from flytekit.core.type_engine import TypeEngine, AsyncTypeTransformer, TypeTransformerFailedError
@@ -339,3 +343,77 @@ def t0(ii: Annotated[MyStr, tf]) -> None:
             print(dynamic_job_spec)
 
     del TypeEngine._REGISTRY[MyStr]
+
+
+# stand-in for prophet.Prophet models
+class Prophet:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+
+def test_annotated_dynamic():
+    # Define a dummy Pydantic dataclass
+    class Config(BaseModel):
+        forecast_horizon: int = 30
+        changepoint_prior_scale: float = 0.05
+        seasonality_mode: str = "multiplicative"
+
+    @task
+    def get_config() -> Config:
+        return Config()
+
+    def dt_function():
+        # Create dummy Prophet models
+        best_models = {
+            "model_1": Prophet(changepoint_prior_scale=0.05),
+            "model_2": Prophet(changepoint_prior_scale=0.1)
+        }
+
+        # Instantiate config
+        cfg = get_config()
+
+        # Create final model dictionary
+        final_model = {
+            "model": best_models,
+            "config": cfg
+        }
+
+        return final_model
+
+    a = Artifact(name="my_model")
+
+    @dynamic
+    def dt1() -> Annotated[dict[str, Union[Config, dict[str, Prophet]]], a]:
+        return dt_function()
+
+    @dynamic
+    def dt_plain() -> dict[str, Union[Config, dict[str, Prophet]]]:
+        return dt_function()
+
+    ss = SerializationSettings(
+        project="test_proj",
+        domain="test_domain",
+        version="abc",
+        image_config=ImageConfig(Image(name="name", fqn="image", tag="name")),
+        env={},
+    )
+
+    with FlyteContextManager.with_context(
+        FlyteContextManager.current_context().with_serialization_settings(ss)
+    ) as ctx:
+        new_exc_state = ctx.execution_state.with_params(mode=ExecutionState.Mode.TASK_EXECUTION)
+        with FlyteContextManager.with_context(ctx.with_execution_state(new_exc_state)):
+            dynamic_job_spec_dt1 = dt1.compile_into_workflow(ctx, dt1._task_function, )
+
+    with FlyteContextManager.with_context(
+        FlyteContextManager.current_context().with_serialization_settings(ss)
+    ) as ctx:
+        new_exc_state = ctx.execution_state.with_params(mode=ExecutionState.Mode.TASK_EXECUTION)
+        with FlyteContextManager.with_context(ctx.with_execution_state(new_exc_state)):
+            dynamic_job_spec_plain = dt_plain.compile_into_workflow(ctx, dt_plain._task_function, )
+
+    assert len(dynamic_job_spec_dt1.nodes) == 1
+    assert len(dynamic_job_spec_plain.nodes) == 1
+
+    assert len(dynamic_job_spec_dt1.outputs[0].binding.map.bindings["model"].map.bindings) == 2
+    assert len(dynamic_job_spec_plain.outputs[0].binding.map.bindings["model"].map.bindings) == 2

tests/flytekit/unit/core/test_generice_idl_type_engine.py

@@ -2858,23 +2858,6 @@ def test_get_underlying_type(t, expected):
     assert get_underlying_type(t) == expected
 
 
-@pytest.mark.parametrize(
-    "t,expected",
-    [
-        (None, (None, None)),
-        (typing.Dict, ()),
-        (typing.Dict[str, str], (str, str)),
-        (
-                Annotated[typing.Dict[str, str], kwtypes(allow_pickle=True)],
-                (typing.Dict[str, str], kwtypes(allow_pickle=True)),
-        ),
-        (typing.Dict[Annotated[str, "a-tag"], int], (Annotated[str, "a-tag"], int)),
-    ],
-)
-def test_dict_get(t, expected):
-    assert DictTransformer.extract_types_or_metadata(t) == expected
-
-
 def test_DataclassTransformer_get_literal_type():
     @dataclass
     class MyDataClassMashumaro(DataClassJsonMixin):

tests/flytekit/unit/core/test_type_engine.py

@@ -2871,13 +2871,13 @@ def test_get_underlying_type(t, expected):
         (typing.Dict[str, str], (str, str)),
         (
                 Annotated[typing.Dict[str, str], kwtypes(allow_pickle=True)],
-                (typing.Dict[str, str], kwtypes(allow_pickle=True)),
+                (str, str),
         ),
         (typing.Dict[Annotated[str, "a-tag"], int], (Annotated[str, "a-tag"], int)),
     ],
 )
 def test_dict_get(t, expected):
-    assert DictTransformer.extract_types_or_metadata(t) == expected
+    assert DictTransformer.extract_types(t) == expected
 
 
 def test_DataclassTransformer_get_literal_type():