feat: integrate TypeID with Pydantic schema handling

This commit adds support for TypeID to work seamlessly with Pydantic's schema generation. By implementing core serialization methods, the goal is to ensure TypeID can be utilized effectively with Pydantic models, enhancing type safety and integration within the application.

Generated-by: aiautocommit

activemodel/__init__.py

@@ -1,2 +1,5 @@
+from . import patches
 from .base_model import BaseModel
+
+# from .field import Field
 from .session_manager import SessionManager, get_engine, get_session, init

activemodel/mixins/typeid.py

@@ -1,5 +1,3 @@
-import uuid
-
 from typeid import TypeID
 
 from activemodel.types.typeid import TypeIDType
@@ -16,7 +14,7 @@ def TypeIDMixin(prefix: str):
     ), f"prefix {prefix} already exists, pick a different one"
 
     class _TypeIDMixin:
-        id: uuid.UUID = Field(
+        id: TypeIDType = Field(
             sa_column=Column(TypeIDType(prefix), primary_key=True),
             default_factory=lambda: TypeID(prefix),
         )

activemodel/patches.py

@@ -0,0 +1,59 @@
+from pydantic import (
+    GetJsonSchemaHandler,
+    WithJsonSchema,
+)
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import CoreSchema, core_schema
+from typeid import TypeID
+
+
+def typeid__get_pydantic_core_schema__(
+    cls, _source_type: type[TypeID], _handler: WithJsonSchema
+) -> CoreSchema:
+    """
+    'Unable to serialize unknown type'
+
+    https://github.com/karma-dev-team/karma-system/blob/ee0c1a06ab2cb7aaca6dc4818312e68c5c623365/app/server/value_objects/steam_id.py#L88
+    https://github.com/hhimanshu/uv-workspaces/blob/main/packages/api/src/_lib/dto/typeid_field.py
+    https://github.com/karma-dev-team/karma-system/blob/ee0c1a06ab2cb7aaca6dc4818312e68c5c623365/app/base/typeid/type_id.py#L14
+    https://github.com/pydantic/pydantic/issues/10060
+    """
+
+    from_uuid_schema = core_schema.chain_schema(
+        [
+            # core_schema.is_instance_schema(TypeID),
+            core_schema.str_schema(),
+            core_schema.no_info_plain_validator_function(TypeID.from_string),
+        ]
+    )
+
+    return core_schema.json_or_python_schema(
+        json_schema=from_uuid_schema,
+        python_schema=core_schema.union_schema(
+            [
+                from_uuid_schema
+                # core_schema.is_instance_schema(TypeID),
+                # core_schema.str_schema(),
+                # core_schema.no_info_plain_validator_function(TypeID.from_string),
+            ]
+        ),
+        serialization=core_schema.plain_serializer_function_ser_schema(
+            lambda x: str(x)
+        ),
+    )
+
+
+def typeid__get_pydantic_json_schema__(
+    cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+) -> JsonSchemaValue:
+    return handler(core_schema.str_schema())
+    return handler(core_schema.uuid_schema())
+
+
+# Monkey-patch the TypeID class
+TypeID.__get_pydantic_core_schema__ = classmethod(typeid__get_pydantic_core_schema__)
+# TypeID.__get_pydantic_json_schema__ = classmethod(typeid__get_pydantic_json_schema__)
+
+# https://github.com/fastapi/fastapi/discussions/10027
+# https://github.com/hhimanshu/uv-workspaces/blob/main/packages/api/src/_lib/dto/typeid_field.py
+# https://github.com/alice-biometrics/petisco/blob/b01ef1b84949d156f73919e126ed77aa8e0b48dd/petisco/base/domain/model/uuid.py#L50

activemodel/types/typeid.py

@@ -5,6 +5,10 @@
 from typing import Optional
 from uuid import UUID
 
+from pydantic import (
+    GetJsonSchemaHandler,
+)
+from pydantic_core import CoreSchema, core_schema
 from typeid import TypeID
 
 from sqlalchemy import types
@@ -87,3 +91,30 @@ def process_result_value(self, value, dialect):
     #         return self
 
     #     return super().coerce_compared_value(op, value)
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> CoreSchema:
+        from_uuid_schema = core_schema.chain_schema(
+            [
+                # core_schema.is_instance_schema(TypeID),
+                core_schema.uuid_schema(),
+                core_schema.no_info_plain_validator_function(TypeID.from_string),
+            ]
+        )
+
+        return core_schema.json_or_python_schema(
+            json_schema=from_uuid_schema,
+            python_schema=core_schema.union_schema(
+                [
+                    from_uuid_schema
+                    # core_schema.is_instance_schema(TypeID),
+                    # core_schema.str_schema(),
+                    # core_schema.no_info_plain_validator_function(TypeID.from_string),
+                ]
+            ),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                lambda x: str(x)
+            ),
+        )

test/typeid_test.py

@@ -1,4 +1,7 @@
+import json
+
 import pytest
+from pydantic import BaseModel as PydanticBaseModel
 from typeid import TypeID
 
 from test.utils import temporary_tables
@@ -26,7 +29,6 @@ class ExampleWithId(BaseModel, TypeIDMixin(TYPEID_PREFIX), table=True):
     pass
 
 
-# the UIDs stored in the DB are NOT the same as the
 def test_get_through_prefixed_uid():
     type_uid = TypeID(prefix=TYPEID_PREFIX)
 
@@ -58,3 +60,24 @@ def test_get_through_plain_uid():
     with temporary_tables():
         record = ExampleWithId.get(type_uid.uuid)
         assert record is None
+
+
+# the wrapped test is probably overkill, but it's protecting against a weird edge case I was running into
+class WrappedExample(PydanticBaseModel):
+    example: ExampleWithId
+
+
+def test_render_typeid():
+    "ensure that pydantic models can render the type id"
+
+    with temporary_tables():
+        example = ExampleWithId().save()
+
+        assert example.model_dump()["id"] == str(example.id)
+        assert json.loads(example.model_dump_json())["id"] == str(example.id)
+
+        wrapped_example = WrappedExample(example=example)
+        assert wrapped_example.model_dump()["example"]["id"] == str(example.id)
+        assert json.loads(wrapped_example.model_dump_json())["example"]["id"] == str(
+            example.id
+        )