implement contains

timgraham · Dec 27, 2024 · f45586c · f45586c
1 parent 6338305
commit f45586c
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 11 deletions.
diff --git a/django_mongodb/features.py b/django_mongodb/features.py
@@ -82,12 +82,13 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "auth_tests.test_views.LoginTest.test_login_session_without_hash_session_key",
         # GenericRelation.value_to_string() assumes integer pk.
         "contenttypes_tests.test_fields.GenericRelationTests.test_value_to_string",
-        # Can contains be implemented?
-        "model_fields_.test_arrayfield.TestQuerying.test_contains",
-        "model_fields_.test_arrayfield.TestQuerying.test_contains_charfield",
+        # contains with expressions/subqueries doesn't work.
         "model_fields_.test_arrayfield.TestQuerying.test_contains_including_expression",
         "model_fields_.test_arrayfield.TestQuerying.test_contains_subquery",
+        # Add FieldGetDbPrepValueMixin to lookups?
+        # https://github.com/django/django/commit/5a36c81f58b8ff45d8dac052343722c54e3fa521#diff-10cd50a6d8761600a1aea96d413184582481a278388b4ad38a788030a21cbd45R5
         "model_fields_.test_arrayfield.TestQuerying.test_enum_lookup",
+        # Unsupported conversion from array to string in $convert
         "model_fields_.test_arrayfield.TestQuerying.test_icontains",
         # Field 'field' expected a number but got Value(1).
         "model_fields_.test_arrayfield.TestQuerying.test_exact_with_expression",

diff --git a/django_mongodb/fields/__init__.py b/django_mongodb/fields/__init__.py
@@ -4,12 +4,7 @@
 from .json import register_json_field
 from .objectid import ObjectIdField
 
-__all__ = [
-    "register_fields",
-    "ArrayField",
-    "ObjectIdAutoField",
-    "ObjectIdField"
-]
+__all__ = ["register_fields", "ArrayField", "ObjectIdAutoField", "ObjectIdField"]
 
 
 def register_fields():

diff --git a/django_mongodb/fields/array.py b/django_mongodb/fields/array.py
@@ -4,12 +4,12 @@
 from django.core import checks, exceptions
 from django.db.models import DecimalField, Field, Func, IntegerField, Transform, Value
 from django.db.models.fields.mixins import CheckFieldDefaultMixin
-from django.db.models.lookups import In
+from django.db.models.lookups import In, Lookup
 from django.utils.translation import gettext_lazy as _
 
 from django_mongodb.forms import SimpleArrayField
 
-from ..query_utils import process_lhs
+from ..query_utils import process_lhs, process_rhs
 from ..utils import prefix_validation_error
 
 __all__ = ["ArrayField"]
@@ -266,6 +266,27 @@ def _rhs_not_none_values(self, rhs):
                 yield True
 
 
+@ArrayField.register_lookup
+class ArrayContains(Lookup):  # ArrayRHSMixin, lookups.DataContains):
+    lookup_name = "contains"
+
+    def as_mql(self, compiler, connection):
+        lhs_mql = process_lhs(self, compiler, connection)
+        value = process_rhs(self, compiler, connection)
+        return {
+            "$gt": [
+                {
+                    "$cond": {
+                        "if": {"$eq": [lhs_mql, None]},
+                        "then": None,
+                        "else": {"$size": {"$setIntersection": [lhs_mql, value]}},
+                    }
+                },
+                0,
+            ]
+        }
+
+
 # @ArrayField.register_lookup
 # class ArrayExact(ArrayRHSMixin, Exact):
 #     pass