🐛 Fix class QueryOperators types tuples

mauro-andre · May 23, 2024 · 1b53e79 · 1b53e79
1 parent f127f02
commit 1b53e79
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 6 deletions.
diff --git a/pyodmongo/models/query_operators.py b/pyodmongo/models/query_operators.py
@@ -61,7 +61,7 @@ class _LogicalOperator(QueryOperator):
     """
 
     operator: str
-    operators: tuple[ComparisonOperator, ...]
+    operators: tuple[Any, ...]
 
     def to_dict(self):
         acu_list = []
@@ -83,7 +83,7 @@ class LogicalOperator(_LogicalOperator):
                           intricate query logic necessary for advanced database operations.
     """
 
-    operators: tuple[ComparisonOperator | _LogicalOperator, ...]
+    operators: tuple[Any, ...]
 
 
 class ElemMatchOperator(QueryOperator):
@@ -102,7 +102,7 @@ class ElemMatchOperator(QueryOperator):
     """
 
     field: Any
-    operators: tuple[ComparisonOperator | _LogicalOperator, ...]
+    operators: tuple[Any, ...]
 
     def to_dict(self):
         elem_match = {}

diff --git a/pyodmongo/queries/query_string.py b/pyodmongo/queries/query_string.py
@@ -1,6 +1,6 @@
 from ..models.db_model import DbModel
 from ..models.db_field_info import DbField
-from ..models.query_operators import LogicalOperator, ComparisonOperator
+from ..models.query_operators import QueryOperator
 from .operators import and_, sort
 from typing import Type
 from datetime import datetime
@@ -69,8 +69,8 @@ def js_regex_to_python(js_regex_str):
 def mount_query_filter(
     Model: Type[DbModel],
     items: dict,
-    initial_comparison_operators: list[ComparisonOperator],
-) -> LogicalOperator:
+    initial_comparison_operators: list[QueryOperator],
+) -> QueryOperator:
     """
     Constructs a MongoDB query filter from a dictionary of conditions and initializes
     additional comparison operators based on the Model's field definitions.

diff --git a/tests/test_queries.py b/tests/test_queries.py
@@ -440,3 +440,39 @@ class MyModel(DbModel):
             }
         }
     }
+
+
+def test_mount_query_filter_with_elem_match():
+    class MyBaseModel(MainBaseModel):
+        attr_1: int
+        attr_2: str
+
+    class MyModel(DbModel):
+        attr_3: str
+        attr_4: list[MyBaseModel]
+        _collection: ClassVar = "my_model"
+
+    dict_input = {"attr_3_eq": "value_3"}
+    elem_match_operator = elem_match(
+        MyBaseModel.attr_1 == "value_1",
+        MyBaseModel.attr_2 == "value_2",
+        field=MyModel.attr_4,
+    )
+    query, _ = mount_query_filter(
+        Model=MyModel,
+        items=dict_input,
+        initial_comparison_operators=[elem_match_operator],
+    )
+    assert query.to_dict() == {
+        "$and": [
+            {
+                "attr_4": {
+                    "$elemMatch": {
+                        "attr_1": {"$eq": "value_1"},
+                        "attr_2": {"$eq": "value_2"},
+                    }
+                }
+            },
+            {"attr_3": {"$eq": "value_3"}},
+        ]
+    }