Support serialization of lists and improve type hints for model binding and result processing

Author: amanmibra

sqlmodel/sql/sqltypes.py

@@ -1,10 +1,12 @@
-from typing import Any, cast
+from typing import Any, List, Type, TypeVar, cast, get_args
 
 from pydantic import BaseModel
 from sqlalchemy import types
 from sqlalchemy.dialects.postgresql import JSONB  # for Postgres JSONB
 from sqlalchemy.engine.interfaces import Dialect
 
+BaseModelType = TypeVar("BaseModelType", bound=BaseModel)
+
 
 class AutoString(types.TypeDecorator):  # type: ignore
     impl = types.String
@@ -24,20 +26,35 @@ class PydanticJSONB(types.TypeDecorator):  # type: ignore
     impl = JSONB  # use JSONB type in Postgres (fallback to JSON for others)
     cache_ok = True  # allow SQLAlchemy to cache results
 
-    def __init__(self, model_class, *args, **kwargs):
+    def __init__(
+        self,
+        model_class: Type[BaseModelType] | Type[list[BaseModelType]],
+        *args,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
         self.model_class = model_class  # Pydantic model class to use
 
-    def process_bind_param(self, value, dialect):
-        # Called when storing to DB: convert Pydantic model to a dict (JSON-serializable)
+    def process_bind_param(self, value: Any, dialect) -> dict | list[dict] | None:  # noqa: ANN401, ARG002, ANN001
         if value is None:
             return None
         if isinstance(value, BaseModel):
-            return value.model_dump()
-        return value  # assume it's already a dict
-
-    def process_result_value(self, value, dialect):
+            return value.model_dump(mode="json")
+        if isinstance(value, list):
+            return [
+                m.model_dump(mode="json") if isinstance(m, BaseModel) else m
+                for m in value
+            ]
+        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
         if value is None:
             return None
-        return self.model_class.parse_obj(value)  # instantiate Pydantic model
+        if isinstance(value, dict):
+            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
+        return value