From e2d95be0843375353d6ec79a230f63d4f874ee0e Mon Sep 17 00:00:00 2001
From: Danil <60117718+damikhai@users.noreply.github.com>
Date: Mon, 26 Feb 2024 23:10:44 +0100
Subject: [PATCH] add possibility of leveraging enumerations in find query
 (#868)

Co-authored-by: Danil Mikhailenko <danil.mikhailenko@t-systems.com>
---
 beanie/odm/utils/encoder.py  |  6 +++++-
 tests/odm/query/test_find.py | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/beanie/odm/utils/encoder.py b/beanie/odm/utils/encoder.py
index 047abbec..6fb6af59 100644
--- a/beanie/odm/utils/encoder.py
+++ b/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]
 
diff --git a/tests/odm/query/test_find.py b/tests/odm/query/test_find.py
index d93f4216..4b120f02 100644
--- a/tests/odm/query/test_find.py
+++ b/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