Implement foo__in filter expressions

ababic · Nov 15, 2018 · 437b846 · 437b846
1 parent 5eb89f0
commit 437b846
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/modelcluster/queryset.py b/modelcluster/queryset.py
@@ -71,6 +71,12 @@ def test_gte(model, attribute_name, value):
     return lambda obj: getattr(obj, attribute_name) >= match_value
 
 
+def test_in(model, attribute_name, value_list):
+    field = model._meta.get_field(attribute_name)
+    match_values = set(field.to_python(val) for val in value_list)
+    return lambda obj: getattr(obj, attribute_name) in match_values
+
+
 FILTER_EXPRESSION_TOKENS = {
     'exact': test_exact,
     'iexact': test_iexact,
@@ -80,6 +86,7 @@ def test_gte(model, attribute_name, value):
     'lte': test_lte,
     'gt': test_gt,
     'gte': test_gte,
+    'in': test_in,
 }
 
 

diff --git a/tests/tests/test_cluster.py b/tests/tests/test_cluster.py
@@ -60,6 +60,9 @@ def test_can_create_cluster(self):
         self.assertEqual(1, beatles.members.filter(name__icontains='carT').count())
         self.assertEqual('Paul McCartney', beatles.members.filter(name__icontains='carT')[0].name)
 
+        self.assertEqual(1, beatles.members.filter(name__in=['Paul McCartney', 'Linda McCartney']).count())
+        self.assertEqual('Paul McCartney', beatles.members.filter(name__in=['Paul McCartney', 'Linda McCartney'])[0].name)
+
         self.assertEqual('Paul McCartney', beatles.members.get(name='Paul McCartney').name)
         self.assertEqual('Paul McCartney', beatles.members.get(name__exact='Paul McCartney').name)
         self.assertEqual('Paul McCartney', beatles.members.get(name__iexact='paul mccartNEY').name)
@@ -69,6 +72,7 @@ def test_can_create_cluster(self):
         self.assertEqual('Paul McCartney', beatles.members.get(name__gte='Paul McCartney').name)
         self.assertEqual('Paul McCartney', beatles.members.get(name__contains='Cart').name)
         self.assertEqual('Paul McCartney', beatles.members.get(name__icontains='carT').name)
+        self.assertEqual('Paul McCartney', beatles.members.get(name__in=['Paul McCartney', 'Linda McCartney']).name)
 
         self.assertRaises(BandMember.DoesNotExist, lambda: beatles.members.get(name='Reginald Dwight'))
         self.assertRaises(BandMember.MultipleObjectsReturned, lambda: beatles.members.get())
@@ -294,6 +298,9 @@ def test_queryset_exclude_filtering(self):
         self.assertEqual(1, beatles.members.exclude(name__icontains='carT').count())
         self.assertEqual('John Lennon', beatles.members.exclude(name__icontains='carT').first().name)
 
+        self.assertEqual(1, beatles.members.exclude(name__in=['Paul McCartney', 'Linda McCartney']).count())
+        self.assertEqual('John Lennon', beatles.members.exclude(name__in=['Paul McCartney', 'Linda McCartney'])[0].name)
+
     def test_prefetch_related(self):
         Band.objects.create(name='The Beatles', members=[
             BandMember(id=1, name='John Lennon'),