Merge pull request #19 from id13/master

Add highlight backend

Author: barseghyanartur

examples/simple/search_indexes/serializers/book.py

@@ -84,6 +84,7 @@ class BookDocumentSimpleSerializer(DocumentSerializer):
 
     # tags = serializers.SerializerMethodField()
     # authors = serializers.SerializerMethodField()
+    highlight = serializers.SerializerMethodField()
 
     class Meta(object):
         """Meta options."""
@@ -103,5 +104,11 @@ class Meta(object):
             'pages',
             'stock_count',
             'tags',
+            'highlight',  # Used in highlight tests
             'null_field',  # Used in testing of `isnull` functional filter.
         )
+
+    def get_highlight(self, obj):
+        if hasattr(obj.meta, 'highlight'):
+            return obj.meta.highlight.__dict__['_d_']
+        return {}

examples/simple/search_indexes/viewsets/book.py

@@ -19,6 +19,7 @@
     OrderingFilterBackend,
     SearchFilterBackend,
     SuggesterFilterBackend,
+    HighlightBackend,
 )
 from django_elasticsearch_dsl_drf.views import BaseDocumentViewSet
 
@@ -47,13 +48,31 @@ class BookDocumentViewSet(BaseDocumentViewSet):
         SearchFilterBackend,
         FacetedSearchFilterBackend,
         SuggesterFilterBackend,
+        HighlightBackend,
     ]
     # Define search fields
     search_fields = (
         'title',
         'description',
         'summary',
     )
+    # Define highlight fields
+    highlight_fields = {
+        'title': {
+            'enabled': True,
+            'options': {
+                'pre_tags': ["<b>"],
+                'post_tags': ["</b>"],
+            }
+        },
+        'summary': {
+            'options': {
+                'fragment_size': 50,
+                'number_of_fragments': 3
+            }
+        },
+        'description': {},
+    }
     # Define filter fields
     filter_fields = {
         'id': {

src/django_elasticsearch_dsl_drf/filter_backends/__init__.py

@@ -15,6 +15,7 @@
 )
 from .search import SearchFilterBackend
 from .suggester import SuggesterFilterBackend
+from .highlight import HighlightBackend
 
 __title__ = 'django_elasticsearch_dsl_drf.filter_backends'
 __author__ = 'Artur Barseghyan <[email protected]>'

src/django_elasticsearch_dsl_drf/filter_backends/highlight.py

@@ -0,0 +1,132 @@
+"""
+Faceted search backend.
+"""
+
+from elasticsearch_dsl.query import Q
+
+from rest_framework.filters import BaseFilterBackend
+
+from six import string_types, iteritems
+
+__title__ = 'django_elasticsearch_dsl_drf.highlight'
+__author__ = 'Artur Barseghyan <[email protected]>'
+__copyright__ = '2017 Artur Barseghyan'
+__license__ = 'GPL 2.0/LGPL 2.1'
+__all__ = ('HighlightBackend',)
+
+
+class HighlightBackend(BaseFilterBackend):
+    """Highlight backend.
+
+    Example:
+
+        >>> from django_elasticsearch_dsl_drf.filter_backends import (
+        >>>     HighlightBackend
+        >>> )
+        >>> from django_elasticsearch_dsl_drf.views import BaseDocumentViewSet
+        >>>
+        >>> # Local article document definition
+        >>> from .documents import ArticleDocument
+        >>>
+        >>> # Local article document serializer
+        >>> from .serializers import ArticleDocumentSerializer
+        >>>
+        >>> class ArticleDocumentView(BaseDocumentViewSet):
+        >>>
+        >>>     document = ArticleDocument
+        >>>     serializer_class = ArticleDocumentSerializer
+        >>>     filter_backends = [HighlightBackend,]
+        >>>     highlight_fields = {
+        >>>         'author.name': {
+        >>>             'enabled': False,
+        >>>             'options': {
+        >>>                 'fragment_size': 150,
+        >>>                 'number_of_fragments': 3
+        >>>             }
+        >>>         }
+        >>>         'title': {
+        >>>             'options': {
+        >>>                 'pre_tags' : ["<em>"],
+        >>>                 'post_tags' : ["</em>"]
+        >>>             },
+        >>>             'enabled': True,
+        >>>         },
+        >>>     }
+
+    Highlight make queries to be more heavy. That's why by default all
+    highlights are disabled and enabled only explicitly either in the filter
+    options (`enabled` set to True) or via query params
+    `?highlight=author.name&highlight=title`.
+    """
+
+    highlight_param = 'highlight'
+
+    @classmethod
+    def prepare_highlight_fields(cls, view):
+        """Prepare faceted search fields.
+
+        Prepares the following structure:
+
+            >>> {
+            >>>     'author.name': {
+            >>>         'enabled': False,
+            >>>         'options': {
+            >>>             'fragment_size': 150,
+            >>>             'number_of_fragments': 3
+            >>>         }
+            >>>     }
+            >>>     'title': {
+            >>>         'options': {
+            >>>             'pre_tags' : ["<em>"],
+            >>>             'post_tags' : ["</em>"]
+            >>>         },
+            >>>         'enabled': True,
+            >>>     },
+            >>> }
+
+        :param view:
+        :type view: rest_framework.viewsets.ReadOnlyModelViewSet
+        :return: Highlight fields options.
+        :rtype: dict
+        """
+        highlight_fields = view.highlight_fields
+
+        for field, options in highlight_fields.items():
+            if 'enabled' not in highlight_fields[field]:
+                highlight_fields[field]['enabled'] = False
+
+            if 'options' not in highlight_fields[field]:
+                highlight_fields[field]['options'] = {}
+
+        return highlight_fields
+
+    def get_highlight_query_params(self, request):
+        """Get highlight query params.
+
+        :param request: Django REST framework request.
+        :type request: rest_framework.request.Request
+        :return: List of search query params.
+        :rtype: list
+        """
+        query_params = request.query_params.copy()
+        return query_params.getlist(self.highlight_param, [])
+
+    def filter_queryset(self, request, queryset, view):
+        """Filter the queryset.
+
+        :param request: Django REST framework request.
+        :param queryset: Base queryset.
+        :param view: View.
+        :type request: rest_framework.request.Request
+        :type queryset: elasticsearch_dsl.search.Search
+        :type view: rest_framework.viewsets.ReadOnlyModelViewSet
+        :return: Updated queryset.
+        :rtype: elasticsearch_dsl.search.Search
+        """
+        highlight_query_params = self.get_highlight_query_params(request)
+        highlight_fields = self.prepare_highlight_fields(view)
+        for __field, __options in highlight_fields.items():
+            if __field in highlight_query_params or __options['enabled']:
+                queryset = queryset.highlight(__field, **__options['options'])
+
+        return queryset

src/django_elasticsearch_dsl_drf/tests/test_highlight.py

@@ -0,0 +1,105 @@
+"""
+Test faceted search backend.
+"""
+
+from __future__ import absolute_import
+
+import unittest
+
+from django.core.management import call_command
+
+from nine.versions import DJANGO_GTE_1_10
+
+import pytest
+
+from rest_framework import status
+
+from books import constants
+import factories
+
+from .base import BaseRestFrameworkTestCase
+
+if DJANGO_GTE_1_10:
+    from django.urls import reverse
+else:
+    from django.core.urlresolvers import reverse
+
+__title__ = 'django_elasticsearch_dsl_drf.tests.test_highlight'
+__author__ = 'Artur Barseghyan <[email protected]>'
+__copyright__ = '2017 Artur Barseghyan'
+__license__ = 'GPL 2.0/LGPL 2.1'
+__all__ = (
+    'TestHighlight',
+)
+
+
+@pytest.mark.django_db
+class TestHighlight(BaseRestFrameworkTestCase):
+    """Test highlight."""
+
+    pytestmark = pytest.mark.django_db
+
+    @classmethod
+    def setUp(cls):
+        cls.books_count = 10
+        cls.books = factories.BookFactory.create_batch(
+            cls.books_count,
+        )
+
+        cls.special_books_count = 10
+        cls.special_books = factories.BookFactory.create_batch(
+            cls.books_count,
+            title='Twenty Thousand Leagues Under the Sea',
+            description="""
+            The title refers to the distance traveled while under the sea and
+            not to a depth, as twenty thousand leagues is over six times the
+            diameter, and nearly twice the circumference of the Earth
+            """,
+            summary="""
+            Margaret Drabble argues that Twenty Thousand Leagues Under the Sea
+            anticipated the ecology movement and shaped the French avant-garde
+            """,
+        )
+        cls.all_books_count = cls.special_books_count + cls.books_count
+
+        call_command('search_index', '--rebuild', '-f')
+
+    def _list_results_with_highlights(self):
+        """List results with facets."""
+        self.authenticate()
+
+        url = reverse('bookdocument-list', kwargs={}) + '?search=twenty'
+        all_highlights_url = url + '&highlight=summary&highlight=description'
+
+        # Make request
+        no_args_response = self.client.get(url, {})
+        self.assertEqual(no_args_response.status_code, status.HTTP_200_OK)
+
+        # Should contain 10 results
+        self.assertEqual(
+            len(no_args_response.data['results']), self.special_books_count)
+
+        for result in no_args_response.data['results']:
+            self.assertEqual(list(result['highlight'].keys()), ['title'])
+
+        # Make request
+        all_highlights_response = self.client.get(all_highlights_url, {})
+        self.assertEqual(
+            all_highlights_response.status_code, status.HTTP_200_OK)
+
+        # Should contain 20 results
+        self.assertEqual(
+            len(all_highlights_response.data['results']),
+            self.special_books_count,
+        )
+        for result in all_highlights_response.data['results']:
+            self.assertEqual(set(['title', 'description', 'summary']),
+                             set(result['highlight'].keys()))
+
+    def test_list_results_with_highlights(self):
+        """Test list results with facets."""
+        return self._list_results_with_highlights()
+
+
+if __name__ == '__main__':
+    unittest.main()

src/django_elasticsearch_dsl_drf/tests/test_search.py

@@ -65,7 +65,7 @@ def setUp(cls):
         cls.switz_cities = factories.CityFactory.create_batch(
             cls.switz_cities_count,
             country=cls.switzerland)
-        cls.all_cities_cound = cls.cities_count + cls.switz_cities_count
+        cls.all_cities_count = cls.cities_count + cls.switz_cities_count
 
         call_command('search_index', '--rebuild', '-f')
 
@@ -102,7 +102,7 @@ def _search_by_nested_field(self, search_term):
         # Should contain 20 results
         response = self.client.get(url, data)
         self.assertEqual(response.status_code, status.HTTP_200_OK)
-        self.assertEqual(len(response.data['results']), self.all_cities_cound)
+        self.assertEqual(len(response.data['results']), self.all_cities_count)
 
         # Should contain only 10 results
         filtered_response = self.client.get(

tox.ini

@@ -14,5 +14,5 @@ deps =
     django110: -r{toxinidir}/examples/requirements/django_1_10.txt
     django111: -r{toxinidir}/examples/requirements/django_1_11.txt
 commands =
-    {envpython} runtests.py
+    {envpython} runtests.py {posargs}
 #    {envpython} examples/simple/manage.py test {posargs:django_elasticsearch_dsl_drf} --settings=settings.testing --traceback -v 3