Dnt create state on init for docs with custom id types (#486)

* fixed

* version: 1.18.0b1

beanie/__init__.py

@@ -27,7 +27,7 @@
 from beanie.odm.views import View
 from beanie.odm.union_doc import UnionDoc
 
-__version__ = "1.18.0b0"
+__version__ = "1.18.0b1"
 __all__ = [
     # ODM
     "Document",

beanie/odm/documents.py

@@ -143,8 +143,6 @@ def _swap_revision(self):
     def __init__(self, *args, **kwargs):
         super(Document, self).__init__(*args, **kwargs)
         self.get_motor_collection()
-        self._save_state()
-        self._swap_revision()
 
     @classmethod
     async def get(

beanie/odm/utils/parsing.py

@@ -33,6 +33,13 @@ def merge_models(left: BaseModel, right: BaseModel) -> None:
             left.__setattr__(k, right_value)
 
 
+def save_state_swap_revision(item: BaseModel):
+    if hasattr(item, "_save_state"):
+        item._save_state()
+    if hasattr(item, "_swap_revision"):
+        item._swap_revision()
+
+
 def parse_obj(
     model: Union[Type[BaseModel], Type["Document"]],
     data: Any,
@@ -76,8 +83,6 @@ def parse_obj(
                 lazy_parse=lazy_parse,
             )  # type: ignore
 
-    # if hasattr(model, "_parse_obj_saving_state"):
-    #     return model._parse_obj_saving_state(data)  # type: ignore
     if (
         lazy_parse
         and hasattr(model, "get_model_type")
@@ -86,4 +91,6 @@ def parse_obj(
         o = model.lazy_parse(data, {"_id"})
         o._saved_state = {"_id": o.id}
         return o
-    return model.parse_obj(data)
+    result = model.parse_obj(data)
+    save_state_swap_revision(result)
+    return result

docs/changelog.md

@@ -2,6 +2,19 @@
 
 Beanie project
 
+## [1.18.0b1] - 2023-02-09
+
+### Fix
+
+- Don't create state on init for docs with custom id types
+
+### Implementation
+
+- Author - [Roman Right](https://github.com/roman-right)
+- PR <https://github.com/roman-right/beanie/pull/486>
+
+[1.18.0b1]: https://pypi.org/project/beanie/1.18.0b1
+
 ## [1.18.0b0] - 2023-01-30
 
 ### Feature

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "beanie"
-version = "1.18.0b0"
+version = "1.18.0b1"
 description = "Asynchronous Python ODM for MongoDB"
 authors = ["Roman <roman-right@protonmail.com>"]
 license = "Apache-2.0"

tests/odm/conftest.py

@@ -69,6 +69,7 @@
     SelfLinked,
     LoopedLinksA,
     LoopedLinksB,
+    DocumentWithTurnedOnStateManagementWithCustomId,
 )
 from tests.odm.views import TestView, TestViewWithLink
 
@@ -223,6 +224,7 @@ async def init(loop, db):
         SelfLinked,
         LoopedLinksA,
         LoopedLinksB,
+        DocumentWithTurnedOnStateManagementWithCustomId,
     ]
     await init_beanie(
         database=db,

tests/odm/models.py

@@ -324,6 +324,15 @@ class Settings:
         use_state_management = True
 
 
+class DocumentWithTurnedOnStateManagementWithCustomId(Document):
+    id: int
+    num_1: int
+    num_2: int
+
+    class Settings:
+        use_state_management = True
+
+
 class DocumentWithTurnedOnReplaceObjects(Document):
     num_1: int
     num_2: int

tests/odm/test_state_management.py

@@ -3,6 +3,7 @@
 
 from beanie import PydanticObjectId, WriteRules
 from beanie.exceptions import StateManagementIsTurnedOff, StateNotSaved
+from beanie.odm.utils.parsing import parse_obj
 from tests.odm.models import (
     DocumentWithTurnedOffStateManagement,
     DocumentWithTurnedOnReplaceObjects,
@@ -13,6 +14,7 @@
     LockWithRevision,
     WindowWithRevision,
     StateAndDecimalFieldModel,
+    DocumentWithTurnedOnStateManagementWithCustomId,
 )
 
 
@@ -37,17 +39,17 @@ def state_without_id():
 
 @pytest.fixture
 def doc_default(state):
-    return DocumentWithTurnedOnStateManagement.parse_obj(state)
+    return parse_obj(DocumentWithTurnedOnStateManagement, state)
 
 
 @pytest.fixture
 def doc_replace(state):
-    return DocumentWithTurnedOnReplaceObjects.parse_obj(state)
+    return parse_obj(DocumentWithTurnedOnReplaceObjects, state)
 
 
 @pytest.fixture
 def doc_previous(state):
-    return DocumentWithTurnedOnSavePrevious.parse_obj(state)
+    return parse_obj(DocumentWithTurnedOnSavePrevious, state)
 
 
 @pytest.fixture
@@ -103,7 +105,7 @@ async def test_parse_object_with_saving_state(self):
             "_id": ObjectId(),
             "internal": InternalDoc(),
         }
-        doc = DocumentWithTurnedOnStateManagement.parse_obj(obj)
+        doc = parse_obj(DocumentWithTurnedOnStateManagement, obj)
         assert doc.get_saved_state() == obj
         assert doc.get_previous_saved_state() is None
 
@@ -136,6 +138,18 @@ async def test_save_state(self):
             }
             assert doc.get_previous_saved_state() is None
 
+        async def test_save_state_with_custom_id_type(self):
+            doc = DocumentWithTurnedOnStateManagementWithCustomId(
+                id=0,
+                num_1=1,
+                num_2=2,
+            )
+            with pytest.raises(StateNotSaved):
+                await doc.save_changes()
+            doc.num_1 = 2
+            with pytest.raises(StateNotSaved):
+                await doc.save_changes()
+
         async def test_save_state_with_previous(self):
             doc = DocumentWithTurnedOnSavePrevious(
                 num_1=1, num_2=2, internal=InternalDoc(num=1, string="s")

tests/test_beanie.py

@@ -2,4 +2,4 @@
 
 
 def test_version():
-    assert __version__ == "1.18.0b0"
+    assert __version__ == "1.18.0b1"