Merge pull request #39 from dabapps/count-and-exists-pair-functions

Count and exists pair functions

README.md

@@ -234,6 +234,8 @@ prepare, project = pairs.combine(
 
 The `pairs.field_display` function takes the field name as its single argument and returns a pair which loads the field from the database, and then projects the result of calling `get_<field>_display` under the key `<field>_display`.
 
+Finally, `pairs.count` and `pairs.has` provide shortcuts to annotate a queryset with the count or existence of related objects, and project these values under the keys `<relationship_name>_count` and `has_<relationship_name>` respectively.
+
 ### `django_readers.specs`: a high-level specification for efficient data querying and projection
 
 This layer is the real magic of `django-readers`: a straightforward way of specifying the shape of your data in order to efficiently select and project a complex tree of related objects.

django_readers/pairs.py

@@ -1,3 +1,4 @@
+from django.db.models import Count
 from django_readers import projectors, qs
 
 
@@ -42,6 +43,24 @@ def field_display(name):
     )
 
 
+def count(name, distinct=True):
+    attr_name = f"{name}_count"
+    return (
+        qs.annotate(**{attr_name: Count(name, distinct=distinct)}),
+        projectors.attr(attr_name),
+    )
+
+
+def has(name, distinct=True):
+    attr_name = f"{name}_count"
+    return (
+        qs.annotate(**{attr_name: Count(name, distinct=distinct)}),
+        projectors.alias(
+            f"has_{name}", projectors.attr(attr_name, transform_value=bool)
+        ),
+    )
+
+
 def filter(*args, **kwargs):
     return prepare_only(qs.filter(*args, **kwargs))
 

tests/test_pairs.py

@@ -587,3 +587,38 @@ def test_pk_list_with_to_attr(self):
         queryset = prepare(Owner.objects.all())
         result = project(queryset.first())
         self.assertEqual(result, {"widgets": [1, 2, 3]})
+
+
+class CountTestCase(TestCase):
+    def test_count(self):
+        owner = Owner.objects.create(name="test owner")
+        Widget.objects.create(name="test 1", owner=owner)
+        Widget.objects.create(name="test 2", owner=owner)
+        Widget.objects.create(name="test 3", owner=owner)
+
+        prepare, project = pairs.count("widget")
+
+        queryset = prepare(Owner.objects.all())
+        result = project(queryset.first())
+        self.assertEqual(result, {"widget_count": 3})
+
+
+class HasTestCase(TestCase):
+    def test_has_false(self):
+        Owner.objects.create(name="test owner")
+
+        prepare, project = pairs.has("widget")
+
+        queryset = prepare(Owner.objects.all())
+        result = project(queryset.first())
+        self.assertEqual(result, {"has_widget": False})
+
+    def test_has_true(self):
+        owner = Owner.objects.create(name="test owner")
+        Widget.objects.create(name="test", owner=owner)
+
+        prepare, project = pairs.has("widget")
+
+        queryset = prepare(Owner.objects.all())
+        result = project(queryset.first())
+        self.assertEqual(result, {"has_widget": True})