diff --git a/atlasq/queryset/transform.py b/atlasq/queryset/transform.py index c4fabe9..09b3cdc 100644 --- a/atlasq/queryset/transform.py +++ b/atlasq/queryset/transform.py @@ -41,12 +41,15 @@ class AtlasTransform: negative_keywords = ["ne", "nin", "not"] exists_keywords = ["exists"] range_keywords = ["gt", "gte", "lt", "lte"] - equals_keywords = ["exact", "iexact", "eq"] + equals_keywords = [] text_keywords = [ "contains", "icontains", "iwholeword", "wholeword", + "exact", + "iexact", + "eq", ] regex_keywords = ["regex", "iregex"] size_keywords = ["size"] @@ -108,12 +111,13 @@ def _exists(self, path: str) -> Dict: return {"exists": {"path": path}} def _range( - self, path: str, value: Union[int, datetime.datetime], keyword: str + self, path: str, value: Union[int, datetime.datetime], keywords: List[str] ) -> Dict: - if keyword not in self.range_keywords: - raise AtlasFieldError( - f"Range search for {path} must be {self.range_keywords}, not {keyword}" - ) + for keyword in keywords: + if keyword not in self.range_keywords: + raise AtlasFieldError( + f"Range search for {path} must be {self.range_keywords}, not {keyword}" + ) if isinstance(value, datetime.datetime): value = value.replace(microsecond=0) elif isinstance(value, int): @@ -122,7 +126,7 @@ def _range( raise AtlasFieldError( f"Range search for {path} must have a value of datetime or integer" ) - return {"range": {"path": path, keyword: value}} + return {"range": {"path": path, **{keyword: value for keyword in keywords}}} def _equals(self, path: str, value: Any) -> Dict: return { @@ -214,7 +218,7 @@ def transform(self) -> Tuple[List[Dict], List[Dict], List[Dict]]: break if keyword in self.range_keywords: - obj = self._range(path, value, keyword) + obj = self._range(path, value, [keyword]) break if keyword in self.equals_keywords: obj = self._equals(path, value) @@ -234,6 +238,8 @@ def transform(self) -> Tuple[List[Dict], List[Dict], List[Dict]]: value_to_check = value if isinstance(value_to_check, (bool, ObjectId)): obj = self._equals(path, value) + elif isinstance(value_to_check, int): + obj = self._range(path, value, ["gte", "lte"]) else: obj = self._text(path, value) diff --git a/tests/queryset/test_transform.py b/tests/queryset/test_transform.py index db18d8e..b9aa850 100644 --- a/tests/queryset/test_transform.py +++ b/tests/queryset/test_transform.py @@ -548,7 +548,7 @@ def test__range_date_valid(self): q = AtlasQ(f=3) t = AtlasTransform(q.query, AtlasIndex("test")) time = datetime.datetime.now() - res = t._range("field", time, "lte") + res = t._range("field", time, ["lte"]) self.assertEqual( res, { @@ -562,7 +562,7 @@ def test__range_date_valid(self): def test__range_integer_valid(self): q = AtlasQ(f=3) t = AtlasTransform(q.query, AtlasIndex("test")) - res = t._range("field", 3, "lte") + res = t._range("field", 3, ["lte"]) self.assertEqual(res, {"range": {"path": "field", "lte": 3}}) def test__range_date_invalid(self): @@ -570,19 +570,19 @@ def test__range_date_invalid(self): t = AtlasTransform(q.query, AtlasIndex("test")) time = datetime.datetime.now() with self.assertRaises(AtlasFieldError): - t._range("field", time, "wat") + t._range("field", time, ["wat"]) def test__range_none(self): q = AtlasQ(f=3) t = AtlasTransform(q.query, AtlasIndex("test")) with self.assertRaises(AtlasFieldError): - t._range("field", None, "lte") + t._range("field", None, ["lte"]) def test__range_string(self): q = AtlasQ(f=3) t = AtlasTransform(q.query, AtlasIndex("test")) with self.assertRaises(AtlasFieldError): - t._range("field", "3", "lte") + t._range("field", "3", ["lte"]) def test__equals(self): q = AtlasQ(f=3) diff --git a/version.py b/version.py index 5a2867d..f406109 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.6.0" +VERSION = "0.7.0"