diff --git a/src/querytyper/mongo/query.py b/src/querytyper/mongo/query.py index 1084ada..70a6130 100644 --- a/src/querytyper/mongo/query.py +++ b/src/querytyper/mongo/query.py @@ -55,11 +55,29 @@ def __or__( class QueryCondition(DictStrAny): - """MongoQuery condition.""" + """Class to represent a single query condition.""" def __init__(self, *args: Any, **kwargs: DictStrAny) -> None: - """Overload init typing.""" - super().__init__(*args, **kwargs) + """ + Initialize a QueryCondition instance. + + It should receive a dict as only argument. + + Example + ------- + ```python + QueryCondition({"field": "value"}) + ``` + + It also overloads dict __init__ typing. + """ + arg = args[0] + if len(args) != 1 or not isinstance(arg, dict): + raise TypeError("QueryCondition must receive only one dict as input.") + if isinstance(arg, dict): + super().__init__(**arg) + for k, v in arg.items(): + self.__setitem__(k, v) def __and__( self, @@ -149,6 +167,13 @@ def __le__( """Overload <= operator.""" return self == {"$lte": other} + def __contains__( + self, + other: T, + ) -> QueryCondition: + """Overload in operator.""" + return regex_query(self.name, re.compile(other)) + def exists( field: Any, @@ -165,6 +190,8 @@ def regex_query( field: Union[QueryField[str], str], regex: re.Pattern, ) -> QueryCondition: - if isinstance(field, QueryField): - return QueryCondition({field.name: {"$regex": regex.pattern}}) - return QueryCondition({field: {"$regex": regex.pattern}}) + return ( + QueryCondition({field.name: {"$regex": regex.pattern}}) + if isinstance(field, QueryField) + else QueryCondition({field: {"$regex": regex.pattern}}) + ) diff --git a/tests/mongo/conftest.py b/tests/mongo/conftest.py index 7d3846b..15190ec 100644 --- a/tests/mongo/conftest.py +++ b/tests/mongo/conftest.py @@ -1,5 +1,6 @@ """Tests fixtures and configurations.""" from typing import Dict, List, Optional + from pydantic import BaseModel from querytyper import MongoFilterMeta diff --git a/tests/mongo/test_pymongo_integration.py b/tests/mongo/test_pymongo_integration.py index 51ce123..01ae691 100644 --- a/tests/mongo/test_pymongo_integration.py +++ b/tests/mongo/test_pymongo_integration.py @@ -1,9 +1,10 @@ """Test querytyper integration with pymongo.""" -import re from typing import Any, Dict + from mongomock import MongoClient -from querytyper import MongoQuery, regex_query +from querytyper import MongoQuery + from .conftest import Dummy, QueryModel @@ -23,10 +24,12 @@ def test_integration_with_pymongo() -> None: for i in range(doc_num) ] ) - # found_doc = collection.find_one(MongoQuery(QueryModel.int_field == 1)._query_dict) found_doc = collection.find_one(MongoQuery(QueryModel.int_field == 1)) assert found_doc is not None found_dummy = Dummy(**found_doc) assert found_dummy.int_field == 1 - found_docs = list(collection.find(MongoQuery(regex_query(QueryModel.str_field, re.compile("test"))))) - assert len(found_docs) == doc_num \ No newline at end of file + query = MongoQuery("test" in QueryModel.str_field) + assert isinstance(query, dict) + # assert query + found_docs = list(collection.find(query)) + assert len(found_docs) == doc_num diff --git a/tests/mongo/test_query.py b/tests/mongo/test_query.py index 07ba4c1..eb02121 100644 --- a/tests/mongo/test_query.py +++ b/tests/mongo/test_query.py @@ -179,3 +179,13 @@ def test_exists_query() -> None: TypeError, match="Field must be a QueryField or str, is not supported." ): exists(1) + + +def test_query_condition_init() -> None: + """Test QueryCondition initializer and TypeErrors.""" + with pytest.raises(TypeError, match="QueryCondition must receive only one dict as input."): + QueryCondition(1) + QueryCondition(1, 2, 3) + condition = QueryCondition({"field": "value"}) + assert "field" in condition + assert condition["field"] == "value"