feat: add contains() operation to QuerySet (#2163)

* Add `contains()` operation to `QuerySet`

* Update the documentation and the changelog

* Trigger CI

* Add docstrings

* Trigger CI

* Update changelog

* Update changelog

* Address PR comments

* Verify the object's model is similar to the queryset's model

* Use `to_db_value()` to get the value

Author: seladb

CHANGELOG.rst

@@ -14,6 +14,7 @@ Changelog
 Added
 ^^^^^
 - ``QuerySet.union()`` — SQL UNION query support for combining results from multiple QuerySets, including support for union across different models, ``union(all=True)`` for duplicates, ``order_by()``, ``limit()``, and ``count()``.
+- ``QuerySet.contains()`` method to check if an object exists in a queryset.
 - Added comprehensive EXPLAIN support for MySQL and PostgreSQL.
 
 Fixed

docs/query.rst

@@ -95,6 +95,19 @@ It's also possible to filter your queries with ``.exclude()``:
 
     await Team.exclude(name__icontains='junior')
 
+You can also check if a specific object exists in a queryset:
+
+.. code-block:: python3
+
+    obj = await Team.filter(name='My Team').first()
+    exists = await Team.all().contains(obj)
+
+Or simply check if any record matches a filter:
+
+.. code-block:: python3
+
+    exists = await Team.filter(name='My Team').exists()
+
 As more interesting case, when you are working with related data, you could also build your
 query around related entities:
 

tests/test_queryset.py

@@ -59,6 +59,30 @@ async def test_exists(db, intfields_data):
     assert not ret
 
 
+@pytest.mark.asyncio
+async def test_contains(db, intfields_data):
+    obj = await IntFields.filter(intnum=10).first()
+    assert await IntFields.all().contains(obj)
+
+    assert await IntFields.filter(intnum__lt=50).contains(obj)
+
+    assert not await IntFields.filter(intnum__gt=50).contains(obj)
+
+
+@pytest.mark.asyncio
+async def test_contains_when_no_pk(db, intfields_data):
+    with pytest.raises(ParamsError, match="The given object does not have a primary key."):
+        await IntFields.all().contains(IntFields(intnum=99))
+
+
+@pytest.mark.asyncio
+async def test_contains_with_wrong_model(db, intfields_data):
+    with pytest.raises(
+        ParamsError, match="The given object is not an instance of the queryset's model."
+    ):
+        await IntFields.all().contains(Tournament(name="test"))
+
+
 @pytest.mark.asyncio
 async def test_limit_count(db, intfields_data):
     assert await IntFields.all().limit(10).count() == 10

tortoise/queryset.py

@@ -822,6 +822,31 @@ def exists(self) -> ExistsQuery:
             use_indexes=self._use_indexes,
         )
 
+    def contains(self, obj: MODEL) -> ContainsQuery:
+        """
+        Check if the QuerySet contains the given instance.
+
+        :param obj: The model instance to check for.
+        :return: True if the QuerySet contains the instance, False otherwise.
+        """
+
+        if not isinstance(obj, self.model):
+            raise ParamsError("The given object is not an instance of the queryset's model.")
+
+        if not obj.pk:
+            raise ParamsError("The given object does not have a primary key.")
+
+        return ContainsQuery(
+            db=self._db,
+            model=self.model,
+            q_objects=self._q_objects,
+            annotations=self._annotations,
+            custom_filters=self._custom_filters,
+            force_indexes=self._force_indexes,
+            use_indexes=self._use_indexes,
+            obj=obj,
+        )
+
     def all(self) -> QuerySet[MODEL]:
         """
         Return the whole QuerySet.
@@ -1473,6 +1498,22 @@ async def _execute(
         return bool(result)
 
 
+class ContainsQuery(ExistsQuery):
+    def __init__(
+        self,
+        obj: MODEL,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._obj = obj
+
+    def _make_query(self) -> None:
+        super()._make_query()
+        pk_field = Field(self.model._meta.db_pk_column)
+        pk_value = self.model._meta.pk.to_db_value(self._obj.pk, self._obj)
+        self.query = self.query.where(pk_field.eq(pk_value))
+
+
 class CountQuery(AwaitableQuery):
     __slots__ = (
         "_limit",