feat: add support computed fields

Author: CatBraaain

beanie/odm/fields.py

@@ -276,6 +276,61 @@ def __deepcopy__(self, memo):
         return self
 
 
+class ExpressionFieldProperty(property):
+    def __init__(
+        self, original_property: property, expression_field: ExpressionField
+    ):
+        self._original = original_property
+        self._expression_field = expression_field
+        super().__init__(
+            original_property.fget,
+            original_property.fset,
+            original_property.fdel,
+            original_property.__doc__,
+        )
+
+    def __getitem__(self, item):
+        return ExpressionField(f"{self._expression_field}.{item}")
+
+    def __getattr__(self, item):
+        return ExpressionField(f"{self._expression_field}.{item}")
+
+    def __hash__(self):
+        return hash(str(self._expression_field))
+
+    def __eq__(self, other):
+        if isinstance(other, ExpressionField):
+            return super(ExpressionField, self._expression_field).__eq__(other)
+        return Eq(field=self, other=other)
+
+    def __gt__(self, other):
+        return GT(field=self._expression_field, other=other)
+
+    def __ge__(self, other):
+        return GTE(field=self._expression_field, other=other)
+
+    def __lt__(self, other):
+        return LT(field=self._expression_field, other=other)
+
+    def __le__(self, other):
+        return LTE(field=self._expression_field, other=other)
+
+    def __ne__(self, other):
+        return NE(field=self._expression_field, other=other)
+
+    def __pos__(self):
+        return self._expression_field, SortDirection.ASCENDING
+
+    def __neg__(self):
+        return self._expression_field, SortDirection.DESCENDING
+
+    def __copy__(self):
+        return self._expression_field
+
+    def __deepcopy__(self, memo):
+        return self._expression_field
+
+
 class DeleteRules(str, Enum):
     DO_NOTHING = "DO_NOTHING"
     DELETE_LINKS = "DELETE_LINKS"

beanie/odm/utils/encoder.py

@@ -27,6 +27,7 @@
 from beanie.odm.utils.pydantic import (
     IS_PYDANTIC_V2,
     IS_PYDANTIC_V2_10,
+    get_model_all_items,
     get_model_fields,
 )
 
@@ -157,7 +158,7 @@ def _iter_model_items(
     ) -> Iterable[Tuple[str, Any]]:
         keep_nulls = self.keep_nulls
         get_model_field = get_model_fields(obj).get
-        for key, value in obj.__iter__():
+        for key, value in get_model_all_items(obj).items():
             field_info = get_model_field(key)
             if field_info is not None:
                 key = field_info.alias or key
@@ -167,7 +168,9 @@ def _iter_model_items(
                 yield key, value
 
     def _should_exclude_field(
-        self, key: str, field_info: Optional[pydantic.fields.FieldInfo]
+        self,
+        key: str,
+        field_info: Any,
     ):
         exclude, include = (
             self.exclude,
@@ -180,7 +183,7 @@ def _should_exclude_field(
         is_pydantic_excluded_field = (
             field_info is not None
             and (
-                field_info.exclude
+                getattr(field_info, "exclude", None)
                 if IS_PYDANTIC_V2
                 else getattr(field_info.field_info, "exclude")
             )

beanie/odm/utils/init.py

@@ -6,6 +6,7 @@
 from beanie.odm.utils.pydantic import (
     IS_PYDANTIC_V2,
     get_extra_field_info,
+    get_field_type,
     get_model_fields,
     parse_model,
 )
@@ -37,6 +38,7 @@
 from beanie.odm.fields import (
     BackLink,
     ExpressionField,
+    ExpressionFieldProperty,
     Link,
     LinkInfo,
     LinkTypes,
@@ -215,9 +217,9 @@ def detect_link(
         :param field: ModelField
         :return: Optional[LinkInfo]
         """
-
-        origin = get_origin(field.annotation)
-        args = get_args(field.annotation)
+        annotation = get_field_type(field)
+        origin = get_origin(annotation)
+        args = get_args(annotation)
         classes = [
             Link,
             BackLink,
@@ -226,8 +228,8 @@ def detect_link(
         for cls in classes:
             # Check if annotation is one of the custom classes
             if (
-                isinstance(field.annotation, _GenericAlias)
-                and field.annotation.__origin__ is cls
+                isinstance(annotation, _GenericAlias)
+                and annotation.__origin__ is cls
             ):
                 if cls is Link:
                     return LinkInfo(
@@ -398,7 +400,14 @@ def init_document_fields(self, cls) -> None:
             cls._link_fields = {}
         for k, v in get_model_fields(cls).items():
             path = v.alias or k
-            setattr(cls, k, ExpressionField(path))
+            attr = getattr(cls, k, None)
+            expression_field = ExpressionField(path)
+            if isinstance(attr, property):
+                setattr(
+                    cls, k, ExpressionFieldProperty(attr, expression_field)
+                )
+            else:
+                setattr(cls, k, expression_field)
 
             link_info = self.detect_link(v, k)
             depth_level = cls.get_settings().max_nesting_depths_per_field.get(

beanie/odm/utils/pydantic.py

@@ -10,6 +10,7 @@
 
 if IS_PYDANTIC_V2:
     from pydantic import TypeAdapter
+    from pydantic.fields import ComputedFieldInfo
 else:
     from pydantic import parse_obj_as
 
@@ -23,18 +24,36 @@ def parse_object_as(object_type: Type, data: Any):
 
 def get_field_type(field):
     if IS_PYDANTIC_V2:
-        return field.annotation
+        if isinstance(field, ComputedFieldInfo):
+            return field.return_type
+        else:
+            return field.annotation
     else:
         return field.outer_type_
 
 
 def get_model_fields(model):
     if IS_PYDANTIC_V2:
-        return model.model_fields
+        return {**model.model_fields, **model.model_computed_fields}
     else:
         return model.__fields__
 
 
+def get_model_all_items(model):
+    if IS_PYDANTIC_V2:
+        return {
+            **dict(model.__iter__()),
+            **{
+                key: getattr(model, key)
+                for key in {
+                    **model.model_computed_fields,
+                }.keys()
+            },
+        }
+    else:
+        return dict(model.__iter__())
+
+
 def parse_model(model_type: Type[BaseModel], data: Any):
     if IS_PYDANTIC_V2:
         return model_type.model_validate(data)

tests/odm/conftest.py

@@ -41,6 +41,7 @@
     DocumentWithBsonBinaryField,
     DocumentWithBsonEncodersFiledsTypes,
     DocumentWithComplexDictKey,
+    DocumentWithComputedField,
     DocumentWithCustomFiledsTypes,
     DocumentWithCustomIdInt,
     DocumentWithCustomIdUUID,
@@ -211,6 +212,7 @@
     BsonRegexDoc,
     NativeRegexDoc,
     DocumentWithExcludedField,
+    DocumentWithComputedField,
 ]
 
 

tests/odm/models.py

@@ -61,7 +61,7 @@
 from beanie.odm.utils.pydantic import IS_PYDANTIC_V2
 
 if IS_PYDANTIC_V2:
-    from pydantic import RootModel, validate_call
+    from pydantic import RootModel, computed_field, validate_call
 
 if sys.version_info >= (3, 10):
 
@@ -912,6 +912,30 @@ class DocumentWithExcludedField(Document):
     excluded_field: Optional[int] = Field(default=None, exclude=True)
 
 
+class DocumentWithComputedField(Document):
+    num: int
+
+    _cached_uuid: Optional[str] = None
+
+    if IS_PYDANTIC_V2:
+
+        @computed_field
+        @property
+        def doubled(self) -> int:
+            return self.num * 2
+
+        @computed_field
+        @property
+        def cacheable_uuid(self) -> str:
+            if self._cached_uuid is None:
+                self._cached_uuid = str(uuid4())
+            return self._cached_uuid
+
+        @cacheable_uuid.setter
+        def cacheable_uuid(self, new: str) -> None:
+            self._cached_uuid = new
+
+
 class ReleaseElemMatch(BaseModel):
     major_ver: int
     minor_ver: int

tests/odm/test_encoder.py

@@ -16,6 +16,7 @@
     DocumentForEncodingTest,
     DocumentForEncodingTestDate,
     DocumentWithComplexDictKey,
+    DocumentWithComputedField,
     DocumentWithDecimalField,
     DocumentWithEnumKeysDict,
     DocumentWithExcludedField,
@@ -146,6 +147,15 @@ async def test_excluded():
     assert "excluded_field" not in encoded_doc
 
 
+@pytest.mark.skipif(not IS_PYDANTIC_V2, reason="Test only for Pydantic v2")
+def test_computed_field():
+    doc = DocumentWithComputedField(num=1)
+    encoded_doc = Encoder().encode(doc)
+    print(doc)
+    print(encoded_doc)
+    assert encoded_doc["doubled"] == 2
+
+
 @pytest.mark.skipif(not IS_PYDANTIC_V2, reason="Test only for Pydantic v2")
 def test_should_encode_pydantic_v2_url_correctly():
     url = AnyUrl("https://example.com")

tests/odm/test_fields.py

@@ -17,6 +17,7 @@
     DocumentTestModel,
     DocumentTestModelIndexFlagsAnnotated,
     DocumentWithBsonEncodersFiledsTypes,
+    DocumentWithComputedField,
     DocumentWithCustomFiledsTypes,
     DocumentWithDeprecatedHiddenField,
     DocumentWithExcludedField,
@@ -124,6 +125,44 @@ async def test_excluded(document):
         assert "excluded_field" not in stored_doc.dict()
 
 
+@pytest.mark.skipif(not IS_PYDANTIC_V2, reason="Test only for Pydantic v2")
+async def test_computed_field():
+    doc = DocumentWithComputedField(num=1)
+    assert doc.doubled == 2
+
+    await doc.insert()
+    stored_doc = await DocumentWithComputedField.get(doc.id)
+    assert stored_doc and stored_doc.doubled == 2
+
+    stored_doc.num = 2
+    assert stored_doc.doubled == 4
+
+    await stored_doc.replace()
+    replaced_doc = await DocumentWithComputedField.get(doc.id)
+    assert replaced_doc and replaced_doc.doubled == 4
+
+
+@pytest.mark.skipif(not IS_PYDANTIC_V2, reason="Test only for Pydantic v2")
+async def test_computed_field_setter():
+    doc = DocumentWithComputedField(num=1)
+    await doc.insert()
+    cached_uui = doc.cacheable_uuid
+    db_raw_data = (
+        await DocumentWithComputedField.get_motor_collection().find_one(
+            {"_id": doc.id}
+        )
+    )
+    assert db_raw_data == {
+        "_id": doc.id,
+        "num": 1,
+        "doubled": 2,
+        "cacheable_uuid": cached_uui,
+    }
+
+    fetched_doc = await DocumentWithComputedField.get(doc.id)
+    assert fetched_doc and fetched_doc.cacheable_uuid != cached_uui
+
+
 async def test_hidden(deprecated_init_beanie):
     document = DocumentWithDeprecatedHiddenField(test_hidden=["abc", "def"])
     await document.insert()