Refactoring to use pydantic-gubbins

Author: nickyoung-github

.github/workflows/python-publish.yml

@@ -0,0 +1,70 @@
+# This workflow will upload a Python Package to PyPI when a release is created
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  release-build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.x"
+
+      - name: Build release distributions
+        run: |
+          # NOTE: put your own distribution build steps here.
+          python -m pip install build
+          python -m build
+
+      - name: Upload distributions
+        uses: actions/upload-artifact@v4
+        with:
+          name: release-dists
+          path: dist/
+
+  pypi-publish:
+    runs-on: ubuntu-latest
+    needs:
+      - release-build
+    permissions:
+      # IMPORTANT: this permission is mandatory for trusted publishing
+      id-token: write
+
+    # Dedicated environments with protections for publishing are strongly recommended.
+    # For more information, see: https://docs.github.com/en/actions/deployment/targeting-different-environments/using-environments-for-deployment#deployment-protection-rules
+    environment:
+      name: pypi
+      # OPTIONAL: uncomment and update to include your PyPI project URL in the deployment status:
+      # url: https://pypi.org/p/YOURPROJECT
+      #
+      # ALTERNATIVE: if your GitHub Release name is the PyPI project version string
+      # ALTERNATIVE: exactly, uncomment the following line instead:
+      # url: https://pypi.org/project/YOURPROJECT/${{ github.event.release.name }}
+
+    steps:
+      - name: Retrieve release distributions
+        uses: actions/download-artifact@v4
+        with:
+          name: release-dists
+          path: dist/
+
+      - name: Publish release distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: dist/

object_model/__init__.py

@@ -1,7 +1,6 @@
 from ._dataclasses import Base, Immutable, NamedPersistable, Persistable
 from ._json import dump, dumps, load, loads
 from ._pydantic import BaseModel, ImmutableModel, NamedPersistableModel, PersistableModel
-from ._typing import Subclass
 from ._descriptors import Id
 
 

object_model/_json.py

@@ -1,14 +1,14 @@
 from dataclasses import MISSING, is_dataclass, fields
-from orjson import loads as __loads, dumps as __dumps
+from orjson import dumps as __dumps
 from pydantic import BaseModel, ConfigDict, TypeAdapter
 from pydantic.alias_generators import to_camel
 from typing import Any
 
 
-from ._type_registry import TYPE_KEY, get_type
+from ._type_registry import get_type
 
 
-__type_adaptors: dict[str, TypeAdapter] = {}
+__type_adaptors: dict[type, TypeAdapter] = {}
 
 
 def get_type_adaptor(typ: type) -> TypeAdapter:
@@ -27,43 +27,32 @@ def get_type_adaptor(typ: type) -> TypeAdapter:
 
 def dump(data: Any) -> dict[str, Any]:
     if isinstance(data, BaseModel):
-        return data.model_dump(include={*data.model_fields_set, TYPE_KEY}, by_alias=True)
+        return data.model_dump(exclude_unset=True, by_alias=True)
     elif is_dataclass(data):
-        flds = [f.name for f in fields(data) if f.default is MISSING or f.default != getattr(data, f.name)]
+        flds = set(f.name for f in fields(data) if f.default is MISSING or f.default != getattr(data, f.name))
         return get_type_adaptor(type(data)).dump_python(data, include=flds)
     else:
         raise RuntimeError("Unsupported type")
 
 
 def dumps(data: Any) -> bytes:
     if isinstance(data, BaseModel):
-        return data.model_dump_json(include={*data.model_fields_set, TYPE_KEY}, by_alias=True).encode("utf-8")
+        return data.model_dump_json(exclude_unset=True, by_alias=True).encode()
     elif is_dataclass(data):
-        flds = [f.name for f in fields(data) if f.default is MISSING or f.default != getattr(data, f.name)]
+        flds = set(f.name for f in fields(data) if f.default is MISSING or f.default != getattr(data, f.name))
         return get_type_adaptor(type(data)).dump_json(data, by_alias=True, include=flds)
     else:
         return __dumps(data)
 
 
-def load(data: dict[str, Any]) -> Any:
-    type_name = data.get(TYPE_KEY)
-    if type_name is None:
-        raise RuntimeError("No type in data")
-
-    typ = get_type(type_name)
+def load(data: dict[str, Any], typ: type | str) -> Any:
+    if isinstance(typ, str):
+        typ = get_type(typ)
 
     return typ.model_validate(data) if issubclass(typ, BaseModel) else get_type_adaptor(typ).validate_python(data)
 
 
-def loads(data: bytes | str, typ: type | str | None = None) -> Any:
-    if typ is None:
-        ret = __loads(data)
-        # SUPER ineffecient ...
-        if TYPE_KEY in ret:
-            typ = ret[TYPE_KEY]
-        else:
-            return ret
-
+def loads(data: bytes | str, typ: type | str) -> Any:
     if isinstance(typ, str):
         typ = get_type(typ)
 

object_model/_type_checking.py

@@ -1,18 +1,12 @@
-from dataclasses import field, is_dataclass
-from datetime import date, datetime
-from pydantic import BaseModel
-from typing import Any, Callable, ClassVar, Literal, Union, get_origin, get_args
+from pydantic._internal._config import ConfigWrapper
+from pydantic_gubbins.typing import FrozenDict, Union
+from typing import Any, Callable, Union as _Union, get_origin, get_args
 
-from ._typing import DiscriminatedUnion, FrozenDict
-from ._type_registry import CLASS_TYPE_KEY, TYPE_KEY, register_type
+from ._type_registry import register_type
 
 
-__base_type_order = {datetime: -2, date: -1, int: -1}
-
-
-def check_type(fld: str, typ: Any) -> Any:
+def check_type(fld: str, typ: Any, immutable_collections: bool) -> Any:
     # Check that we have no non-serialisable or ambiguously serialisable types
-    # Also, rewrite a couple of types to avoid common problems
 
     if typ in (object, Any, Callable):
         raise TypeError(f"{typ} is not a persistable type for {fld}")
@@ -22,50 +16,33 @@ def check_type(fld: str, typ: Any) -> Any:
 
     if not args:
         if origin in (dict, list, set, tuple):
-            raise TypeError(f"Cannot use untyped collection for {field}")
+            raise TypeError(f"Cannot use untyped collection for {fld}")
 
     for arg in args:
-        check_type(fld, arg)
-
-    if origin is set:
-        return frozenset[args]
-    elif origin is list:
-        return tuple[args + (...,)]
-    elif origin is dict:
-        return FrozenDict[args]
-    elif origin is Union:
-        # Re-order the args of unions so that e.g. datetime comes before str
-
-        object_types = tuple(t for t in args if issubclass(t, BaseModel) or is_dataclass(t))
-        base_types = set(args).difference(object_types) if object_types else args
-        base_types = tuple(sorted(base_types, key=lambda x: __base_type_order.get(x, 0)))
+        check_type(fld, arg, immutable_collections)
 
-        # If we have e.g. Union[date, MyClass, MyOtherClass] we need to use a discriminated union for the
-        # classes, so we need Union[date, Annotated[Union[MyClass, MyOtherClass], Field(discriminator=TYPE_KEY)]
+    if origin is _Union:
+        return Union[args]
 
-        if len(object_types) > 1:
-            return Union[base_types + DiscriminatedUnion[object_types]] if base_types else\
-                DiscriminatedUnion[object_types]
+    if immutable_collections:
+        if origin is set:
+            return frozenset[args]
+        elif origin is list:
+            return tuple[args + (...,)]
+        elif origin is dict:
+            return FrozenDict[args]
 
     return typ
 
 
 class TypeCheckMixin:
     def __new__(cls, cls_name: str, bases: tuple[type[Any], ...], namespace: dict[str, Any], **kwargs):
-        annotations = namespace.setdefault("__annotations__", {})
-
-        for name, typ in annotations.items():
-            annotations[name] = check_type(name, typ)
-
-        registered_name = annotations.get(TYPE_KEY)
-
-        if not registered_name:
-            registered_name = cls_name
-            annotations[TYPE_KEY] = Literal[registered_name]
-            namespace[TYPE_KEY] = field(default_factory=lambda: registered_name, init=False)
+        model_config = ConfigWrapper.for_model(bases, namespace, kwargs)
+        immutable_collections = model_config.frozen if model_config else False
+        annotations = namespace.get("__annotations__", {})
 
-        annotations[CLASS_TYPE_KEY] = ClassVar[str]
-        namespace[CLASS_TYPE_KEY] = registered_name
+        for name, typ in namespace.setdefault("__annotations__", {}).items():
+            annotations[name] = check_type(name, typ, immutable_collections)
 
         ret = super().__new__(cls, cls_name, bases, namespace, **kwargs)
         register_type(ret)

object_model/_type_registry.py

@@ -1,7 +1,5 @@
 import importlib.metadata as md
-
-TYPE_KEY = "t_"
-CLASS_TYPE_KEY = "t__"
+from pydantic_gubbins.typing import get_type_name
 
 
 class __TypeRegistry:
@@ -34,10 +32,7 @@ def is_temporary_type(self, type_name: str) -> bool:
         return is_temporary
 
     def register_type(self, typ: type):
-        type_name = getattr(typ, CLASS_TYPE_KEY, None)
-        if type_name is None:
-            raise RuntimeError(f"{typ} is missing attribute {TYPE_KEY}")
-
+        type_name = get_type_name(typ)
         if type_name in self.__object_store.names:
             return
 
@@ -53,9 +48,7 @@ def get_type(type_name: str) -> type:
 
 def is_temporary_type(typ: str | type) -> bool:
     if isinstance(typ, type):
-        type_name = getattr(typ, CLASS_TYPE_KEY, None)
-        if type_name is None:
-            raise RuntimeError(f"{typ} is missing attribute {TYPE_KEY}")
+        type_name = get_type_name(typ)
     else:
         type_name = typ
 

object_model/_typing.py

@@ -31,7 +31,7 @@ def message(cls) -> str:
 
 class WrongStoreError(ObjectStoreError):
     def __init__(self, object_type: str, object_id: bytes):
-        message = f"""Attempting to save object {object_type}:{object_id.decode("utf-8")}
+        message = f"""Attempting to save object {object_type}:{object_id.decode()}
                       in a different store to the one from which it was loaded"""
         super().__init__(message)
 

object_model/store/exception.py

@@ -5,13 +5,14 @@
 from orjson import loads
 from platform import system, uname
 from pydantic import BaseModel
+from pydantic_gubbins.typing import get_type_name
 from typing import Iterable
 
 from . import ObjectResult
 from .exception import NotFoundError
 from .persistable import ImmutableMixin, ObjectRecord, PersistableMixin
 from .._json import schema
-from .._type_registry import CLASS_TYPE_KEY, is_temporary_type
+from .._type_registry import is_temporary_type
 
 
 def _get_user_name():
@@ -145,7 +146,7 @@ def write(self, obj: PersistableMixin, as_of_effective_time: bool = False) -> Fu
         return ret
 
     def register_type(self, typ: type[PersistableMixin]):
-        self.register_schema(RegisterSchemaRequest(name=getattr(typ, CLASS_TYPE_KEY), json_schema=schema(typ)))
+        self.register_schema(RegisterSchemaRequest(name=get_type_name(typ), json_schema=schema(typ)))
 
     def register_schema(self, request: RegisterSchemaRequest):
         defs = request.json_schema.pop("$defs", {})

object_model/store/object_store.py

@@ -4,12 +4,12 @@
 from datetime import datetime
 from functools import cached_property
 from hashlib import sha3_512
+from pydantic_gubbins.typing import get_type_name
 from uuid import UUID
 
 from .object_record import ObjectRecord
 from .._descriptors import Id
 from .._json import dumps, loads
-from .._type_registry import CLASS_TYPE_KEY, TYPE_KEY
 
 
 class UseDerived:
@@ -51,11 +51,11 @@ def id(self) -> tuple[type, tuple[str, ...]]:
 
     @property
     def object_type(self) -> str:
-        return getattr(self, TYPE_KEY)
+        return get_type_name(type(self))
 
     @property
     def object_id_type(self) -> str:
-        return getattr(self.id[0], CLASS_TYPE_KEY)
+        return get_type_name(self.id[0])
 
     @property
     def object_id(self) -> bytes:
@@ -81,7 +81,7 @@ def make_id(cls, *args, **kwargs) -> tuple[str, bytes]:
             else:
                 raise ValueError(f"Missing ID field {name}")
 
-        return getattr(id_type, CLASS_TYPE_KEY), dumps(ret)
+        return get_type_name(id_type), dumps(ret)
 
     @classmethod
     def _check_persistable_class(cls):