add possibility of leveraging enumerations in find query (#868)

Co-authored-by: Danil Mikhailenko <danil.mikhailenko@t-systems.com>

beanie/odm/utils/encoder.py

@@ -7,6 +7,7 @@
 import pathlib
 import re
 import uuid
+from enum import Enum
 from typing import (
     Any,
     Callable,
@@ -131,7 +132,10 @@ def encode(self, obj: Any) -> Any:
             items = self._iter_model_items(obj)
             return {key: self.encode(value) for key, value in items}
         if isinstance(obj, Mapping):
-            return {str(key): self.encode(value) for key, value in obj.items()}
+            return {
+                key if isinstance(key, Enum) else str(key): self.encode(value)
+                for key, value in obj.items()
+            }
         if isinstance(obj, Iterable):
             return [self.encode(value) for value in obj]
 

tests/odm/query/test_find.py

@@ -1,4 +1,5 @@
 import datetime
+from enum import Enum
 
 import pytest
 from pydantic import BaseModel
@@ -399,3 +400,19 @@ def test_find_clone():
         ("string", SortDirection.ASCENDING),
     ]
     assert new_q.limit_number == 10
+
+
+async def test_find_many_with_enum_in_query(preset_documents):
+    class TestEnum(str, Enum):
+        INTEGER = Sample.integer
+        SAMPLE_NESTED_OPTIONAL = Sample.nested.optional
+        CONST = "const"
+        CONST_VALUE = "TEST"
+
+    filter_query = {
+        TestEnum.INTEGER: {"$gt": 1},
+        TestEnum.SAMPLE_NESTED_OPTIONAL: {"$type": "null"},
+        TestEnum.CONST: TestEnum.CONST_VALUE,
+    }
+    result = await Sample.find_many(filter_query).to_list()
+    assert len(result) == 2