Skip to content

Commit

Permalink
add possibility of leveraging enumerations in find query
Browse files Browse the repository at this point in the history
  • Loading branch information
Danil Mikhailenko committed Feb 19, 2024
1 parent 72b35f9 commit 0d60c9a
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
6 changes: 5 additions & 1 deletion beanie/odm/utils/encoder.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import pathlib
import re
import uuid
from enum import Enum
from typing import (
Any,
Callable,
Expand Down Expand Up @@ -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]

Expand Down
17 changes: 17 additions & 0 deletions tests/odm/query/test_find.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import datetime
from enum import Enum

import pytest
from pydantic import BaseModel
Expand Down Expand Up @@ -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

0 comments on commit 0d60c9a

Please sign in to comment.