diff --git a/README.md b/README.md index c0ff9b1..7cc71d6 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,7 @@ from shylock import ShylockAioArangoDBBackend from shylock import configure as configure_shylock from arangodantic import ASCENDING, DocumentModel, EdgeModel, configure +from arangodantic.operators import Operators # Define models @@ -127,9 +128,28 @@ async def main(): async for found_company in cursor: print(f"Company: {found_company.company_id}") - # Supported operators include: "==", "!=", "<", "<=", ">", ">=" + # Supported operators include (see arangodantic.operators for more): + # Operators.EQ (==) + # Operators.NE (!=) + # Operators.LT (<) + # Operators.LTE (<=) + # Operators.GT (>) + # Operators.GTE (>=) + # Operators.IN + # Operators.LIKE + # Operators.NOT_LIKE + # Operators.REG_MATCH (=~) + # Operators.NOT_REG_MATCH (!~) + # Operators.NOT (!) + # Operators.NOT_IN + # Operators.ALL_IN + # Operators.ANY_IN + # Operators.NONE_IN + # Operators.ALL_(EQ, NE, LT, LTE, GT, GTE) (ALL ==, ALL !=, ALL <, etc.) + # Operators.ANY_(EQ, NE, LT, LTE, GT, GTE) (ANY ==, ANY !=, ANY <, etc.) + # Operators.NONE_(EQ, NE, LT, LTE, GT, GTE) (NONE ==, NONE !=, NONE <, etc.) found_company = await Company.find_one( - {"owner.last_name": "Doe", "_id": {"!=": company}} + {"owner.last_name": "Doe", "_id": {Operators.NE: company}} ) print(f"Found the company {found_company.key_}") diff --git a/arangodantic/operators.py b/arangodantic/operators.py new file mode 100644 index 0000000..f729c1d --- /dev/null +++ b/arangodantic/operators.py @@ -0,0 +1,89 @@ +from enum import Enum + + +class Operators(str, Enum): + """ + Enum class for different operators. + + See: + https://www.arangodb.com/docs/stable/aql/operators.html#comparison-operators + https://www.arangodb.com/docs/stable/aql/operators.html#array-comparison-operators + for more details. + """ + + EQ = "==" + NE = "!=" + LT = "<" + LTE = "<=" + GT = ">" + GTE = ">=" + IN = "IN" + NOT_IN = "NOT IN" + LIKE = "LIKE" + NOT_LIKE = "NOT LIKE" + REG_MATCH = "=~" + NOT_REG_MATCH = "!~" + NOT = "NOT" + ALL_IN = "ALL IN" + NONE_IN = "NONE IN" + ANY_IN = "ANY IN" + # These array comparison operators requires an extra operator + # which will compare each individual array element. + ANY_EQ = "ANY ==" + ANY_NE = "ANY !=" + ANY_LT = "ANY <" + ANY_LTE = "ANY <=" + ANY_GT = "ANY >" + ANY_GTE = "ANY >=" + ALL_EQ = "==" + ALL_NE = "!=" + ALL_LT = "ALL <" + ALL_LTE = "ALL <=" + ALL_GT = "ALL >" + ALL_GTE = "ALL >=" + NONE_EQ = "NONE ==" + NONE_NE = "NONE !=" + NONE_LT = "NONE <" + NONE_LTE = "NONE <=" + NONE_GT = "NONE >" + NONE_GTE = "NONE >=" + + +# List of supported operators mapped to a-z string representations that can be +# used safely in the names of bind_vars in AQL +comparison_operators = { + Operators.EQ: "eq", + Operators.NE: "ne", + Operators.LT: "lt", + Operators.LTE: "lte", + Operators.GT: "gt", + Operators.GTE: "gte", + Operators.IN: "in", + Operators.NOT_IN: "not_in", + Operators.LIKE: "like", + Operators.NOT_LIKE: "not_like", + Operators.REG_MATCH: "reg_match", + Operators.NOT_REG_MATCH: "not_reg_match", + Operators.NOT: "not", + Operators.ALL_IN: "all_in", + Operators.NONE_IN: "none_in", + Operators.ANY_IN: "any_in", + Operators.ANY_EQ: "any_eq", + Operators.ANY_NE: "any_ne", + Operators.ANY_LT: "any_lt", + Operators.ANY_LTE: "any_lte", + Operators.ANY_GT: "any_gt", + Operators.ANY_GTE: "any_gte", + Operators.ALL_EQ: "all_eq", + Operators.ALL_NE: "all_ne", + Operators.ALL_LT: "all_lt", + Operators.ALL_LTE: "all_lte", + Operators.ALL_GT: "all_gt", + Operators.ALL_GTE: "all_gte", + Operators.NONE_EQ: "none_eq", + Operators.NONE_NE: "none_ne", + Operators.NONE_LT: "none_lt", + Operators.NONE_LTE: "none_lte", + Operators.NONE_GT: "none_gt", + Operators.NONE_GTE: "none_gte", +} diff --git a/arangodantic/tests/conftest.py b/arangodantic/tests/conftest.py index 4fa59d7..8c17a86 100644 --- a/arangodantic/tests/conftest.py +++ b/arangodantic/tests/conftest.py @@ -1,6 +1,6 @@ import random import string -from typing import Optional +from typing import Any, Dict, Optional from uuid import uuid4 import pydantic @@ -52,6 +52,7 @@ class Identity(DocumentModel): """Dummy identity Arangodantic model.""" name: str = "" + data: Optional[Dict[str, Any]] = None class SubModel(pydantic.BaseModel): diff --git a/arangodantic/tests/test_model.py b/arangodantic/tests/test_model.py index aca2e70..bbba532 100644 --- a/arangodantic/tests/test_model.py +++ b/arangodantic/tests/test_model.py @@ -13,6 +13,7 @@ MultipleModelsFoundError, UniqueConstraintError, ) +from arangodantic.operators import Operators from arangodantic.tests.conftest import ExtendedIdentity, Identity, Link, SubModel from arangodantic.utils import SortTypes @@ -167,8 +168,8 @@ async def test_find(identity_collection): async def test_find_with_comparisons(identity_collection): i_a = Identity(name="a") i_a2 = Identity(name="a") - i_b = Identity(name="b") - i_c = Identity(name="c") + i_b = Identity(name="b", data={"aliases": ["foo", "bar"]}) + i_c = Identity(name="c", data={"aliases": ["bar", "baz"]}) await gather(i_a.save(), i_a2.save(), i_b.save(), i_c.save()) @@ -178,44 +179,76 @@ async def test_find_with_comparisons(identity_collection): async for i in cursor: assert i.name == "a" - cursor = await (Identity.find({"name": {"<": "a"}}, count=True)) + cursor = await (Identity.find({"name": {Operators.LT: "a"}}, count=True)) async with cursor: assert len(cursor) == 0 - cursor = await (Identity.find({"name": {"<=": "a"}}, count=True)) + cursor = await (Identity.find({"name": {Operators.LTE: "a"}}, count=True)) async with cursor: assert len(cursor) == 2 async for i in cursor: assert i.name == "a" - cursor = await (Identity.find({"name": {">": "c"}}, count=True)) + cursor = await (Identity.find({"name": {Operators.GT: "c"}}, count=True)) async with cursor: assert len(cursor) == 0 - cursor = await (Identity.find({"name": {">": "b"}}, count=True)) + cursor = await (Identity.find({"name": {Operators.GT: "b"}}, count=True)) async with cursor: assert len(cursor) == 1 async for i in cursor: assert i.name == "c" - cursor = await (Identity.find({"name": {">=": "b"}}, count=True)) + cursor = await (Identity.find({"name": {Operators.GTE: "b"}}, count=True)) async with cursor: assert len(cursor) == 2 async for i in cursor: assert i.name in {"b", "c"} - cursor = await (Identity.find({"name": {">": "a", "<": "c"}}, count=True)) + cursor = await ( + Identity.find({"name": {Operators.GT: "a", Operators.LT: "c"}}, count=True) + ) async with cursor: assert len(cursor) == 1 async for i in cursor: assert i.name == "b" - cursor = await (Identity.find({"name": "a", "_id": {"!=": i_a}}, count=True)) + cursor = await ( + Identity.find({"name": "a", "_id": {Operators.NE: i_a}}, count=True) + ) async with cursor: assert len(cursor) == 1 async for i in cursor: assert i.id_ == i_a2.id_ + cursor = await (Identity.find({"name": {Operators.IN: ["b", "c"]}}, count=True)) + async with cursor: + assert len(cursor) == 2 + async for i in cursor: + assert i.name in {"b", "c"} + + cursor = await (Identity.find({"name": {Operators.NOT_IN: ["a"]}}, count=True)) + async with cursor: + assert len(cursor) == 2 + async for i in cursor: + assert i.name in {"b", "c"} + + cursor = await ( + Identity.find({"data.aliases": {Operators.ALL_IN: ["foo", "bar"]}}, count=True) + ) + async with cursor: + assert len(cursor) == 1 + async for i in cursor: + assert i.name == "b" + + cursor = await ( + Identity.find({"data.aliases": {Operators.NONE_IN: ["foo"]}}, count=True) + ) + async with cursor: + assert len(cursor) == 1 + async for i in cursor: + assert i.name == "c" + @pytest.mark.parametrize( "bad_str", diff --git a/arangodantic/tests/test_operators.py b/arangodantic/tests/test_operators.py new file mode 100644 index 0000000..be9b68c --- /dev/null +++ b/arangodantic/tests/test_operators.py @@ -0,0 +1,89 @@ +from typing import Any, List, Union + +import pytest + +from arangodantic.operators import Operators +from arangodantic.tests.conftest import Identity + + +@pytest.mark.parametrize( + "data, operator, query, expected_length", + [ + (0, Operators.EQ, None, 0), + (1, Operators.GT, 0, 1), + (True, Operators.NE, None, 1), + (45, Operators.LTE, "yikes!", 1), + (65, Operators.NE, "65", 1), + (65, Operators.EQ, 65, 1), + (1.23, Operators.GT, 1.32, 0), + (1.5, Operators.IN, [2, 3, 1.5], 1), + ("foo", Operators.IN, None, 0), + (42, Operators.NOT_IN, [17, 40, 50], 1), + ("abc", Operators.EQ, "abc", 1), + ("abc", Operators.EQ, "ABC", 0), + ("foo", Operators.LIKE, "f%", 1), + ("foo", Operators.NOT_LIKE, "f%", 0), + ("foo", Operators.REG_MATCH, "^f[o].$", 1), + ("foo", Operators.NOT_REG_MATCH, "[a-z]+bar$", 1), + ], +) +@pytest.mark.asyncio +async def test_comparison_operators( + identity_collection, + data: Any, + operator: Operators, + query: Any, + expected_length: int, +): + # See https://www.arangodb.com/docs/stable/aql/operators.html#comparison-operators + # for the test cases. + i_a = Identity(name="a", data={"value": data}) + await i_a.save() + + cursor = await Identity.find({"data.value": {operator: query}}, count=True) + async with cursor: + assert len(cursor) == expected_length + + +@pytest.mark.parametrize( + "data, operator, query, expected_length", + [ + ([1, 2, 3], Operators.ALL_IN, [2, 3, 4], 0), + ([1, 2, 3], Operators.ALL_IN, [1, 2, 3], 1), + ([1, 2, 3], Operators.NONE_IN, [3], 0), + ([1, 2, 3], Operators.NONE_IN, [23, 42], 1), + ([1, 2, 3], Operators.ANY_IN, [4, 5, 6], 0), + ([1, 2, 3], Operators.ANY_IN, [1, 42], 1), + ([1, 2, 3], Operators.ANY_EQ, 2, 1), + ([1, 2, 3], Operators.ANY_EQ, 4, 0), + ([1, 2, 3], Operators.ANY_GT, 0, 1), + ([1, 2, 3], Operators.ANY_LTE, 1, 1), + ([1, 2, 3], Operators.ANY_LTE, 1, 1), + ([1, 2, 3], Operators.NONE_LT, 99, 0), + ([1, 2, 3], Operators.NONE_GT, 10, 1), + ([1, 2, 3], Operators.ALL_GT, 2, 0), + ([1, 2, 3], Operators.ALL_GT, 0, 1), + ([1, 2, 3], Operators.ALL_GTE, 3, 0), + (["foo", "bar"], Operators.ALL_NE, "moo", 1), + (["foo", "bar"], Operators.NONE_EQ, "bar", 0), + (["foo", "bar"], Operators.ANY_EQ, "foo", 1), + ], +) +@pytest.mark.asyncio +async def test_array_comparison_operators( + identity_collection, + data: List[Union[int, str]], + operator: Operators, + query: Union[List[int], int], + expected_length: int, +): + # See + # https://www.arangodb.com/docs/stable/aql/operators.html#array-comparison-operators + # for the test cases. + i_a = Identity(name="a", data={"list": data}) + + await i_a.save() + + cursor = await Identity.find({"data.list": {operator: query}}, count=True) + async with cursor: + assert len(cursor) == expected_length diff --git a/arangodantic/utils.py b/arangodantic/utils.py index c3ae7b5..b7802db 100644 --- a/arangodantic/utils.py +++ b/arangodantic/utils.py @@ -2,6 +2,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple from arangodantic.directions import DIRECTIONS +from arangodantic.operators import comparison_operators FilterTypes = Optional[Dict[str, Any]] SortTypes = Optional[Iterable[Tuple[str, str]]] @@ -37,17 +38,6 @@ def build_filters( filter_list = [] bind_vars = {} - # List of supported operators mapped to a-z string representations that can be - # used safely in the names of bind_vars in AQL - comparison_operators = { - "<": "lt", - "<=": "lte", - ">": "gt", - ">=": "gte", - "!=": "ne", - "==": "eq", - } - if filters: for i, (field, expr) in enumerate(filters.items()): if not isinstance(expr, Dict): diff --git a/examples/readme_example.py b/examples/readme_example.py index e9f000c..5d0fdaf 100644 --- a/examples/readme_example.py +++ b/examples/readme_example.py @@ -8,6 +8,7 @@ from shylock import configure as configure_shylock from arangodantic import ASCENDING, DocumentModel, EdgeModel, configure +from arangodantic.operators import Operators # Define models @@ -88,9 +89,28 @@ async def main(): async for found_company in cursor: print(f"Company: {found_company.company_id}") - # Supported operators include: "==", "!=", "<", "<=", ">", ">=" + # Supported operators include (see arangodantic.operators for more): + # Operators.EQ (==) + # Operators.NE (!=) + # Operators.LT (<) + # Operators.LTE (<=) + # Operators.GT (>) + # Operators.GTE (>=) + # Operators.IN + # Operators.LIKE + # Operators.NOT_LIKE + # Operators.REG_MATCH (=~) + # Operators.NOT_REG_MATCH (!~) + # Operators.NOT (!) + # Operators.NOT_IN + # Operators.ALL_IN + # Operators.ANY_IN + # Operators.NONE_IN + # Operators.ALL_(EQ, NE, LT, LTE, GT, GTE) (ALL ==, ALL !=, ALL <, etc.) + # Operators.ANY_(EQ, NE, LT, LTE, GT, GTE) (ANY ==, ANY !=, ANY <, etc.) + # Operators.NONE_(EQ, NE, LT, LTE, GT, GTE) (NONE ==, NONE !=, NONE <, etc.) found_company = await Company.find_one( - {"owner.last_name": "Doe", "_id": {"!=": company}} + {"owner.last_name": "Doe", "_id": {Operators.NE: company}} ) print(f"Found the company {found_company.key_}")