Add extensible Indexable method to get items for indexing #32 (#33)

parasolr/indexing.py

@@ -69,6 +69,18 @@ def index_item_type(cls):
         in index manage command. """
         return cls._meta.verbose_name
 
+    @classmethod
+    def items_to_index(cls):
+        """Get all items to be indexed for a single class of Indexable
+        content. Subclasses can override this method to return a custom
+        iterable, e.g. a Django `QuerySet` that takes advantage of
+        prefetching. By default, returns all Django objects for a model.
+        Raises NotImplementedError if that fails."""
+        try:
+            return cls.objects.all()
+        except AttributeError:
+            raise NotImplementedError
+
     def index_id(self):
         """Solr identifier. By default, combines :meth:`index item_type`
         and :attr:`id` with :attr:ID_SEPARATOR`."""

parasolr/management/commands/index.py

@@ -116,7 +116,19 @@ def handle(self, *args, **kwargs):
             # calculate total to index across all indexables for current mode
             for name, model in self.indexables.items():
                 if self.options['index'] in [name, 'all']:
-                    total_to_index += model.objects.count()
+                    # possibly inefficient to generate the list just
+                    # for a count; should be ok for django queryset implementation,
+                    # hopefully not too bad for other cases
+                    items = model.items_to_index()
+                    if items:
+                        try:
+                            # try count, since it's more effecient for
+                            # django querysets
+                            total_to_index += items.count()
+                        except TypeError:
+                            # if count errors because we have a list,
+                            # use len
+                            total_to_index += len(items)
 
         # initialize progressbar if requested and indexing more than 5 items
         progbar = None
@@ -135,7 +147,7 @@ def handle(self, *args, **kwargs):
             for name, model in self.indexables.items():
                 if self.options['index'] in [name, 'all']:
                     # index in chunks and update progress bar
-                    count += self.index(model.objects.all(), progbar=progbar)
+                    count += self.index(model.items_to_index(), progbar=progbar)
 
         if progbar:
             progbar.finish()

parasolr/tests/test_indexing.py

@@ -114,3 +114,23 @@ def test_index_items__queryset(self, mocksolr):
         mockqueryset = MagicMock(spec=QuerySet)
         Indexable.index_items(mockqueryset)
         mockqueryset.iterator.assert_called_with()
+
+    def test_items_to_index(self, mocksolr):
+        # assumes django model manager interface by default
+
+        # simple object with objects.all interface
+        simple_items_to_index = SimpleIndexable.items_to_index()
+        assert len(simple_items_to_index) == 5
+        assert isinstance(simple_items_to_index[0], SimpleIndexable)
+
+        # model-ish object
+        model_items_to_index = ModelIndexable.items_to_index()
+        assert len(model_items_to_index) == 1
+        assert isinstance(model_items_to_index[0], ModelIndexable)
+
+        class NonModelIndexable(Indexable):
+            pass
+
+        # raises not implemented if objects.all fails
+        with pytest.raises(NotImplementedError):
+            NonModelIndexable.items_to_index()