Support custom sqltype with MutablePydanticBaseModel

Author: wonderbeyond

.vscode/settings.json

@@ -1,5 +1,6 @@
 {
     "cSpell.words": [
-        "autouse"
+        "autouse",
+        "sqltype"
     ]
 }

sqlalchemy_nested_mutable/mutable.py

@@ -45,15 +45,20 @@ class PydanticType(sa.types.TypeDecorator, TypeEngine[_P]):
         """
         Inspired by https://gist.github.com/imankulov/4051b7805ad737ace7d8de3d3f934d6b
         """
+        cache_ok = True
         impl = sa.types.JSON
 
-        def __init__(self, pydantic_type: type[_P]):
+        def __init__(self, pydantic_type: type[_P], sqltype: TypeEngine[_T] = None):
             super().__init__()
             self.pydantic_type = pydantic_type
+            self.sqltype = sqltype
 
         def load_dialect_impl(self, dialect):
             from sqlalchemy.dialects.postgresql import JSONB
 
+            if self.sqltype is not None:
+                return dialect.type_descriptor(self.sqltype)
+
             if dialect.name == "postgresql":
                 return dialect.type_descriptor(JSONB())
             return dialect.type_descriptor(sa.JSON())
@@ -79,8 +84,8 @@ def dict(self, *args, **kwargs):
             return res
 
         @classmethod
-        def as_mutable(cls, /) -> TypeEngine[Self]:
-            return super().as_mutable(PydanticType(cls))
+        def as_mutable(cls, sqltype: TypeEngine[_T] = None) -> TypeEngine[Self]:
+            return super().as_mutable(PydanticType(cls, sqltype))
 elif not TYPE_CHECKING:
     class PydanticType:
         def __new__(cls, *a, **k):

tests/test_custom_sqltype.py

@@ -0,0 +1,55 @@
+from typing import Optional, List
+
+import pytest
+from sqlalchemy.dialects.postgresql import JSON, JSONB
+from sqlalchemy_nested_mutable._compat import pydantic
+import sqlalchemy as sa
+from sqlalchemy.orm import (
+    DeclarativeBase,
+    Mapped,
+    mapped_column,
+)
+
+from sqlalchemy_nested_mutable import MutablePydanticBaseModel
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class Addresses(MutablePydanticBaseModel):
+    class AddressItem(pydantic.BaseModel):
+        street: str
+        city: str
+        area: Optional[str]
+
+    work: List[AddressItem] = []
+    home: List[AddressItem] = []
+
+
+class User(Base):
+    __tablename__ = "user_account"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[str] = mapped_column(sa.String(30))
+    addresses_default: Mapped[Optional[Addresses]] = mapped_column(Addresses.as_mutable())
+    addresses_json: Mapped[Optional[Addresses]] = mapped_column(Addresses.as_mutable(JSON))
+    addresses_jsonb: Mapped[Optional[Addresses]] = mapped_column(Addresses.as_mutable(JSONB))
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _with_tables(session):
+    Base.metadata.create_all(session.bind)
+    yield
+    session.execute(sa.text("""
+    DROP TABLE user_account CASCADE;
+    """))
+    session.commit()
+
+
+def test_mutable_pydantic_type(session):
+    session.add(User(name="foo"))
+    session.commit()
+    assert session.scalar(sa.select(sa.func.pg_typeof(User.addresses_default))) == "jsonb"
+    assert session.scalar(sa.select(sa.func.pg_typeof(User.addresses_json))) == "json"
+    assert session.scalar(sa.select(sa.func.pg_typeof(User.addresses_jsonb))) == "jsonb"