Add support for Indexed custom pydantic fields (#754)

BeanieODM · Oct 24, 2023 · 46683ce · 46683ce
1 parent 57dd7e7
commit 46683ce
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/beanie/odm/fields.py b/beanie/odm/fields.py
@@ -86,6 +86,12 @@ def __new__(cls, *args, **kwargs):
             def __get_pydantic_core_schema__(
                 cls, _source_type: Any, _handler: GetCoreSchemaHandler
             ) -> core_schema.CoreSchema:
+                custom_type = getattr(
+                    typ, "__get_pydantic_core_schema__", None
+                )
+                if custom_type is not None:
+                    return custom_type(_source_type, _handler)
+
                 return core_schema.no_info_after_validator_function(
                     lambda v: v,
                     simple_ser_schema(typ.__name__),

diff --git a/tests/odm/documents/test_init.py b/tests/odm/documents/test_init.py
@@ -6,6 +6,7 @@
 from beanie.exceptions import CollectionWasNotInitialized
 from beanie.odm.utils.projection import get_projection
 from tests.odm.models import (
+    Color,
     DocumentTestModel,
     DocumentTestModelStringImport,
     DocumentTestModelWithComplexIndex,
@@ -289,3 +290,20 @@ async def test_merge_indexes():
 
 async def test_custom_init():
     assert DocumentWithCustomInit.s == "TEST2"
+
+
+async def test_index_on_custom_types(db):
+    class Sample1(Document):
+        name: Indexed(Color, unique=True)
+
+        class Settings:
+            name = "sample"
+
+    await db.drop_collection("sample")
+
+    await init_beanie(
+        database=db,
+        document_models=[Sample1],
+    )
+
+    await db.drop_collection("sample")