new: Support ManyToManyFields in model_attr lookups

Terr · acdha · commit 2a450d81432a · 2016-05-16T09:46:15.000-04:00
Thanks to @Terr for the patch
diff --git a/AUTHORS b/AUTHORS
@@ -111,3 +111,4 @@ Thanks to
     * Claude Paroz (@claudep) for Django 1.9 support
     * Chris Brooke (@chrisbrooke) for patching around a backwards-incompatible change in ElasticSearch 2
     * Gilad Beeri (@giladbeeri) for adding retries when updating a backend
+    * Arjen Verstoep (@terr) for a patch that allows attribute lookups through Django ManyToManyField relationships
diff --git a/haystack/fields.py b/haystack/fields.py
@@ -80,40 +80,80 @@ def prepare(self, obj):
         if self.use_template:
             return self.prepare_template(obj)
         elif self.model_attr is not None:
-            # Check for `__` in the field for looking through the relation.
-            attrs = self.model_attr.split('__')
-            current_object = obj
-
-            for attr in attrs:
-                if not hasattr(current_object, attr):
-                    raise SearchFieldError("The model '%s' does not have a model_attr '%s'." % (repr(current_object), attr))
-
-                current_object = getattr(current_object, attr, None)
-
-                if current_object is None:
-                    if self.has_default():
-                        current_object = self._default
-                        # Fall out of the loop, given any further attempts at
-                        # accesses will fail miserably.
-                        break
-                    elif self.null:
-                        current_object = None
-                        # Fall out of the loop, given any further attempts at
-                        # accesses will fail miserably.
-                        break
-                    else:
-                        raise SearchFieldError("The model '%s' combined with model_attr '%s' returned None, but doesn't allow a default or null value." % (repr(obj), self.model_attr))
+            attrs = self.split_model_attr_lookups()
+            current_objects = [obj]
 
-            if callable(current_object):
-                return current_object()
+            values = self.resolve_attributes_lookup(current_objects, attrs)
 
-            return current_object
+            if len(values) == 1:
+                return values[0]
+            else:
+                return values
 
         if self.has_default():
             return self.default
         else:
             return None
 
+    def resolve_attributes_lookup(self, current_objects, attributes):
+        """
+        Recursive method that looks, for one or more objects, for an attribute that can be multiple
+        objects (relations) deep.
+        """
+        values = []
+
+        for current_object in current_objects:
+            if not hasattr(current_object, attributes[0]):
+                raise SearchFieldError(
+                    "The model '%s' does not have a model_attr '%s'." % (repr(current_object), attributes[0])
+                )
+
+            if len(attributes) > 1:
+                current_objects_in_attr = self.get_iterable_objects(getattr(current_object, attributes[0]))
+
+                return self.resolve_attributes_lookup(current_objects_in_attr, attributes[1:])
+
+            current_object = getattr(current_object, attributes[0])
+
+            if current_object is None:
+                if self.has_default():
+                    current_object = self._default
+                elif self.null:
+                    current_object = None
+                else:
+                    raise SearchFieldError(
+                        "The model '%s' combined with model_attr '%s' returned None, but doesn't allow "
+                        "a default or null value." % (repr(current_object), self.model_attr)
+                    )
+
+            if callable(current_object):
+                values.append(current_object())
+            else:
+                values.append(current_object)
+
+        return values
+
+    def split_model_attr_lookups(self):
+        """Returns list of nested attributes for looking through the relation."""
+        return self.model_attr.split('__')
+
+    @classmethod
+    def get_iterable_objects(cls, current_objects):
+        """
+        Returns iterable of objects that contain data. For example, resolves Django ManyToMany relationship
+        so the attributes of the related models can then be accessed.
+        """
+        if current_objects is None:
+            return []
+
+        if hasattr(current_objects, 'all'):
+            # i.e, Django ManyToMany relationships
+            current_objects = current_objects.all()
+        elif not hasattr(current_objects, '__iter__'):
+            current_objects = [current_objects]
+
+        return current_objects
+
     def prepare_template(self, obj):
         """
         Flattens an object for indexing.
diff --git a/test_haystack/core/models.py b/test_haystack/core/models.py
@@ -79,3 +79,22 @@ class ScoreMockModel(models.Model):
 
     def __unicode__(self):
         return self.score
+
+
+class ManyToManyLeftSideModel(models.Model):
+    related_models = models.ManyToManyField('ManyToManyRightSideModel')
+
+
+class ManyToManyRightSideModel(models.Model):
+    name = models.CharField(max_length=32, default='Default name')
+
+    def __unicode__(self):
+        return self.name
+
+
+class OneToManyLeftSideModel(models.Model):
+    pass
+
+
+class OneToManyRightSideModel(models.Model):
+    left_side = models.ForeignKey(OneToManyLeftSideModel, related_name='right_side')
diff --git a/test_haystack/test_fields.py b/test_haystack/test_fields.py
@@ -5,13 +5,87 @@
 import datetime
 from decimal import Decimal
 
+from mock import Mock
+
 from django.template import TemplateDoesNotExist
 from django.test import TestCase
-from test_haystack.core.models import MockModel, MockTag
+from test_haystack.core.models import MockModel, MockTag, ManyToManyLeftSideModel, ManyToManyRightSideModel, \
+    OneToManyLeftSideModel, OneToManyRightSideModel
 
 from haystack.fields import *
 
 
+class SearchFieldTestCase(TestCase):
+    def test_get_iterable_objects_with_none(self):
+        self.assertEqual([], SearchField.get_iterable_objects(None))
+
+    def test_get_iterable_objects_with_single_non_iterable_object(self):
+        obj = object()
+        expected = [obj]
+
+        self.assertEqual(expected, SearchField.get_iterable_objects(obj))
+
+    def test_get_iterable_objects_with_list_stays_the_same(self):
+        objects = [object(), object()]
+
+        self.assertIs(objects, SearchField.get_iterable_objects(objects))
+
+    def test_get_iterable_objects_with_django_manytomany_rel(self):
+        left_model = ManyToManyLeftSideModel.objects.create()
+        right_model_1 = ManyToManyRightSideModel.objects.create(name='Right side 1')
+        right_model_2 = ManyToManyRightSideModel.objects.create()
+        left_model.related_models.add(right_model_1)
+        left_model.related_models.add(right_model_2)
+
+        result = SearchField.get_iterable_objects(left_model.related_models)
+
+        self.assertTrue(right_model_1 in result)
+        self.assertTrue(right_model_2 in result)
+
+    def test_get_iterable_objects_with_django_onetomany_rel(self):
+        left_model = OneToManyLeftSideModel.objects.create()
+        right_model_1 = OneToManyRightSideModel.objects.create(left_side=left_model)
+        right_model_2 = OneToManyRightSideModel.objects.create(left_side=left_model)
+
+        result = SearchField.get_iterable_objects(left_model.right_side)
+
+        self.assertTrue(right_model_1 in result)
+        self.assertTrue(right_model_2 in result)
+
+    def test_resolve_attributes_lookup_with_field_that_points_to_none(self):
+        related = Mock(spec=['none_field'], none_field=None)
+        obj = Mock(spec=['related'], related=[related])
+
+        field = SearchField(null=False)
+
+        self.assertRaises(SearchFieldError, field.resolve_attributes_lookup, [obj], ['related', 'none_field'])
+
+    def test_resolve_attributes_lookup_with_field_that_points_to_none_but_is_allowed_to_be_null(self):
+        related = Mock(spec=['none_field'], none_field=None)
+        obj = Mock(spec=['related'], related=[related])
+
+        field = SearchField(null=True)
+
+        self.assertEqual([None], field.resolve_attributes_lookup([obj], ['related', 'none_field']))
+
+    def test_resolve_attributes_lookup_with_field_that_points_to_none_but_has_default(self):
+        related = Mock(spec=['none_field'], none_field=None)
+        obj = Mock(spec=['related'], related=[related])
+
+        field = SearchField(default='Default value')
+
+        self.assertEqual(['Default value'], field.resolve_attributes_lookup([obj], ['related', 'none_field']))
+
+    def test_resolve_attributes_lookup_with_deep_relationship(self):
+        related_lvl_2 = Mock(spec=['value'], value=1)
+        related = Mock(spec=['related'], related=[related_lvl_2, related_lvl_2])
+        obj = Mock(spec=['related'], related=[related])
+
+        field = SearchField()
+
+        self.assertEqual([1, 1], field.resolve_attributes_lookup([obj], ['related', 'related', 'value']))
+
+
 class CharFieldTestCase(TestCase):
     def test_init(self):
         try:
diff --git a/test_haystack/test_indexes.py b/test_haystack/test_indexes.py
@@ -8,7 +8,8 @@
 
 from django.test import TestCase
 from django.utils.six.moves import queue
-from test_haystack.core.models import AFifthMockModel, AThirdMockModel, MockModel
+from test_haystack.core.models import (AFifthMockModel, AThirdMockModel, MockModel, ManyToManyLeftSideModel,
+                                       ManyToManyRightSideModel)
 
 from haystack import connection_router, connections, indexes
 from haystack.exceptions import SearchFieldError
@@ -134,6 +135,14 @@ class MROFieldsSearchChild(MROFieldsSearchIndexA, MROFieldsSearchIndexB):
     pass
 
 
+class ModelWithManyToManyFieldAndAttributeLookupSearchIndex(indexes.SearchIndex, indexes.Indexable):
+    text = indexes.CharField(document=True)
+    related_models = indexes.MultiValueField(model_attr='related_models__name')
+
+    def get_model(self):
+        return ManyToManyLeftSideModel
+
+
 class SearchIndexTestCase(TestCase):
     fixtures = ['base_data']
 
@@ -650,3 +659,27 @@ def test_float_integer_fields(self):
         self.assertTrue(isinstance(self.yabmsi.fields['average_delay'], indexes.FloatField))
         self.assertEqual(self.yabmsi.fields['average_delay'].null, False)
         self.assertEqual(self.yabmsi.fields['average_delay'].index_fieldname, 'average_delay')
+
+
+class ModelWithManyToManyFieldAndAttributeLookupSearchIndexTestCase(TestCase):
+    def test_full_prepare(self):
+        index = ModelWithManyToManyFieldAndAttributeLookupSearchIndex()
+
+        left_model = ManyToManyLeftSideModel.objects.create()
+        right_model_1 = ManyToManyRightSideModel.objects.create(name='Right side 1')
+        right_model_2 = ManyToManyRightSideModel.objects.create()
+        left_model.related_models.add(right_model_1)
+        left_model.related_models.add(right_model_2)
+
+        result = index.full_prepare(left_model)
+
+        self.assertDictEqual(
+            result,
+            {
+                'django_ct': 'core.manytomanyleftsidemodel',
+                'django_id': '1',
+                'text': None,
+                'id': 'core.manytomanyleftsidemodel.1',
+                'related_models': ['Right side 1', 'Default name'],
+            }
+        )
