Django Haystack backend support for Solr's Collapsing query parser and Expand component

collapsed.py

# encoding: utf-8
"""
Experimental SearchQuerySet exposing Solr's Collapse filter and Expand component

See https://cwiki.apache.org/confluence/display/solr/Collapse+and+Expand+Results

Usage::
    sqs = sqs.collapse('item_grouping', sort='"score DESC, wdl_id ASC"')
    sqs = sqs.expand(rows=2, fq='django_ct:core.file')
    
Note that _process_results has some non-portable code to handle inconsistencies
in the response types for numeric item IDs in JSON.
"""
from __future__ import absolute_import, division, print_function

from haystack.backends.solr_backend import SolrEngine, SolrSearchBackend, SolrSearchQuery
from haystack.query import SearchQuerySet


class CollapsedSearchQuerySet(SearchQuerySet):
    def __init__(self, *args, **kwargs):
        super(CollapsedSearchQuerySet, self).__init__(*args, **kwargs)

    def collapse(self, field_name, **kwargs):
        # To implement collapsing we'll need to add a filter-query like this:
        #   &fq={!collapse field=item_grouping sort="django_ct desc, score desc"}

        # FIXME: handle escaping, etc.
        fq = ['field=%s' % field_name]

        for k, v in kwargs.items():
            fq.append('%s=%s' % (k, v))

        return self.narrow('{!collapse %s}' % ' '.join(fq))

    def expand(self, **kwargs):
        clone = self._clone()
        assert isinstance(clone.query, CollapsedSearchQuery)
        clone.query.add_expand(**kwargs)
        return clone


class CollapsedSearchQuery(SolrSearchQuery):
    def __init__(self, *args, **kwargs):
        super(CollapsedSearchQuery, self).__init__(*args, **kwargs)
        self.expand = None

    def add_expand(self, **kwargs):
        self.expand = kwargs

    def build_params(self, *args, **kwargs):
        params = super(CollapsedSearchQuery, self).build_params(*args, **kwargs)

        if self.expand is not None:
            params['expand'] = 'true'
            for k, v in self.expand.items():
                params['expand.%s' % k] = v

        return params


class CollapsedSolrSearchBackend(SolrSearchBackend):
    def _process_results(self, raw_results, *args, **kwargs):
        res = super(CollapsedSolrSearchBackend, self)._process_results(raw_results, *args, **kwargs)

        if 'expanded' in raw_results.raw_response:
            # The keys should be numeric (WDL ID or Item Grouping) but they'll be strings in the raw
            # JSON response since JSON inherits JavaScript's lack of integer safety. We'll convert:
            expanded = {int(k): v for k, v in raw_results.raw_response['expanded'].items()}

            for i in res['results']:
                extra = expanded.get(i.wdl_id) or expanded.get(i.item_grouping)
                setattr(i, 'expanded_results', extra)

        return res


class CollapsedSolrEngine(SolrEngine):
    backend = CollapsedSolrSearchBackend
    query = CollapsedSearchQuery