Support dict[str, BaseModel]

Author: amanmibra

sqlmodel/sql/sqltypes.py

@@ -1,4 +1,4 @@
-from typing import Any, List, Type, TypeVar, cast, get_args
+from typing import Any, Dict, List, Type, TypeVar, cast, get_args
 
 from pydantic import BaseModel
 from sqlalchemy import types
@@ -28,7 +28,9 @@ class PydanticJSONB(types.TypeDecorator):  # type: ignore
 
     def __init__(
         self,
-        model_class: Type[BaseModelType] | Type[list[BaseModelType]],
+        model_class: Type[BaseModelType]
+        | Type[list[BaseModelType]]
+        | Type[Dict[str, BaseModelType]],
         *args,
         **kwargs,
     ):
@@ -45,15 +47,23 @@ def process_bind_param(self, value: Any, dialect) -> dict | list[dict] | None:
                 m.model_dump(mode="json") if isinstance(m, BaseModel) else m
                 for m in value
             ]
+        if isinstance(value, dict):
+            return {
+                k: v.model_dump(mode="json") if isinstance(v, BaseModel) else v
+                for k, v in value.items()
+            }
         return value
 
     def process_result_value(
         self, value: Any, dialect
-    ) -> BaseModelType | List[BaseModelType] | None:  # noqa: ANN401, ARG002, ANN001
-        # Called when loading from DB: convert dict to Pydantic model instance
+    ) -> BaseModelType | List[BaseModelType] | Dict[str, BaseModelType] | None:  # noqa: ANN401, ARG002, ANN001
         if value is None:
             return None
         if isinstance(value, dict):
+            # If model_class is dict, handle key-value pairs
+            if isinstance(self.model_class, dict):
+                return {k: self.model_class.model_validate(v) for k, v in value.items()}
+            # Regular case: the whole dict represents a single model
             return self.model_class.model_validate(value)  # type: ignore
         if isinstance(value, list):
             return [get_args(self.model_class)[0].model_validate(v) for v in value]  # type: ignore