Increase operator support

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_}")
 

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",
+}

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):

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",

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

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):