fix: Pydantic "exclude" option is not working #756 (#1154)

* fix: respect exclude option in Pydantic fields

Author: CatBraaain

beanie/odm/utils/dump.py

@@ -16,9 +16,12 @@ def get_dict(
         exclude = set()
     if document.id is None:
         exclude.add("_id")
-    if not document.get_settings().use_revision:
-        exclude.add("revision_id")
-    encoder = Encoder(exclude=exclude, to_db=to_db, keep_nulls=keep_nulls)
+    include = set()
+    if document.get_settings().use_revision:
+        include.add("revision_id")
+    encoder = Encoder(
+        exclude=exclude, include=include, to_db=to_db, keep_nulls=keep_nulls
+    )
     return encoder.encode(document)
 
 

beanie/odm/utils/encoder.py

@@ -83,6 +83,7 @@ class Encoder:
     """
 
     exclude: Container[str] = frozenset()
+    include: Container[str] = frozenset()
     custom_encoders: Mapping[type, SingleArgCallable] = dc.field(
         default_factory=dict
     )
@@ -154,15 +155,39 @@ def encode(self, obj: Any) -> Any:
     def _iter_model_items(
         self, obj: pydantic.BaseModel
     ) -> Iterable[Tuple[str, Any]]:
-        exclude, keep_nulls = self.exclude, self.keep_nulls
+        keep_nulls = self.keep_nulls
         get_model_field = get_model_fields(obj).get
         for key, value in obj.__iter__():
             field_info = get_model_field(key)
             if field_info is not None:
                 key = field_info.alias or key
-            if key not in exclude and (value is not None or keep_nulls):
+            if not self._should_exclude_field(key, field_info) and (
+                value is not None or keep_nulls
+            ):
                 yield key, value
 
+    def _should_exclude_field(
+        self, key: str, field_info: Optional[pydantic.fields.FieldInfo]
+    ):
+        exclude, include = (
+            self.exclude,
+            self.include,
+        )
+
+        if key in include:
+            return False
+
+        is_pydantic_excluded_field = (
+            field_info is not None
+            and (
+                field_info.exclude
+                if IS_PYDANTIC_V2
+                else getattr(field_info.field_info, "exclude")
+            )
+            is True
+        )
+        return key in exclude or is_pydantic_excluded_field
+
 
 def _get_encoder(
     obj: Any, custom_encoders: Mapping[type, SingleArgCallable]

tests/odm/conftest.py

@@ -48,6 +48,7 @@
     DocumentWithDecimalField,
     DocumentWithDeprecatedHiddenField,
     DocumentWithEnumKeysDict,
+    DocumentWithExcludedField,
     DocumentWithExtras,
     DocumentWithHttpUrlField,
     DocumentWithIndexedObjectId,
@@ -209,6 +210,7 @@
     LongSelfLink,
     BsonRegexDoc,
     NativeRegexDoc,
+    DocumentWithExcludedField,
 ]
 
 

tests/odm/models.py

@@ -165,7 +165,7 @@ class DocumentTestModel(Document):
     test_int: int
     test_doc: SubDocument
     test_str: str
-    test_list: List[SubDocument] = Field(exclude=True)
+    test_list: List[SubDocument]
 
     class Settings:
         use_cache = True
@@ -268,9 +268,11 @@ class Settings:
 
 class DocumentWithDeprecatedHiddenField(Document):
     if IS_PYDANTIC_V2:
-        test_hidden: List[str] = Field(json_schema_extra={"hidden": True})
+        test_hidden: Optional[List[str]] = Field(
+            default=None, json_schema_extra={"hidden": True}
+        )
     else:
-        test_hidden: List[str] = Field(hidden=True)
+        test_hidden: Optional[List[str]] = Field(default=None, hidden=True)
 
 
 class DocumentWithCustomIdUUID(Document):
@@ -568,7 +570,7 @@ class House(Document):
     roof: Optional[Link[Roof]] = None
     yards: Optional[List[Link[Yard]]] = None
     height: Indexed(int) = 2
-    name: Indexed(str) = Field(exclude=True)
+    name: Indexed(str)
 
     if IS_PYDANTIC_V2:
         model_config = ConfigDict(
@@ -905,6 +907,11 @@ class Settings:
         use_state_management = True
 
 
+class DocumentWithExcludedField(Document):
+    included_field: int
+    excluded_field: Optional[int] = Field(default=None, exclude=True)
+
+
 class ReleaseElemMatch(BaseModel):
     major_ver: int
     minor_ver: int

tests/odm/test_encoder.py

@@ -18,6 +18,7 @@
     DocumentWithComplexDictKey,
     DocumentWithDecimalField,
     DocumentWithEnumKeysDict,
+    DocumentWithExcludedField,
     DocumentWithHttpUrlField,
     DocumentWithKeepNullsFalse,
     DocumentWithStringField,
@@ -138,6 +139,13 @@ def test_keep_nulls_false():
     assert encoded_doc == {"m": {"i": 10}}
 
 
+async def test_excluded():
+    doc = DocumentWithExcludedField(included_field=1, excluded_field=2)
+    encoded_doc = Encoder().encode(doc)
+    assert "included_field" in encoded_doc
+    assert "excluded_field" not in encoded_doc
+
+
 @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

@@ -19,6 +19,7 @@
     DocumentWithBsonEncodersFiledsTypes,
     DocumentWithCustomFiledsTypes,
     DocumentWithDeprecatedHiddenField,
+    DocumentWithExcludedField,
     Sample,
 )
 
@@ -111,11 +112,16 @@ async def test_custom_filed_types():
 
 
 async def test_excluded(document):
-    document = await DocumentTestModel.find_one()
+    doc = DocumentWithExcludedField(included_field=1, excluded_field=2)
+    await doc.insert()
+    stored_doc = await DocumentWithExcludedField.get(doc.id)
+    assert stored_doc is not None
     if IS_PYDANTIC_V2:
-        assert "test_list" not in document.model_dump()
+        assert "included_field" in stored_doc.model_dump()
+        assert "excluded_field" not in stored_doc.model_dump()
     else:
-        assert "test_list" not in document.dict()
+        assert "included_field" in stored_doc.dict()
+        assert "excluded_field" not in stored_doc.dict()
 
 
 async def test_hidden(deprecated_init_beanie):