from typing import Dict, List
from parasolr.query.queryset import SolrQuerySet
[docs]class AliasedSolrQuerySet(SolrQuerySet):
"""Extension of :class:`~parasolr.query.queryset.SolrQuerySet`
with support for aliasing Solr fields to more readable versions
for use in code. To use, extend this class and define a
dictionary of :attr:`field_aliases` with the same syntax you would
when calling :meth:`only`. Those field aliases will be set
as the default initial value for :attr:`field_list`, and aliases
can be used in all extended methods.
"""
#: map of application-specific, readable field names
#: to actual solr fields (i.e. if using dynamic field types)
field_aliases = {}
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# set default field list based on field_aliases
self.field_list = [
"%s:%s" % (key, value) for key, value in self.field_aliases.items()
]
# generate reverse lookup for updating facets & highlights
self.reverse_aliases = {val: key for key, val in self.field_aliases.items()}
[docs] def _unalias_args(self, *args):
"""convert alias name to solr field for list of args"""
return [self.field_aliases.get(arg, arg) for arg in args]
[docs] def _unalias_kwargs(self, **kwargs):
"""convert alias name to solr field for keys in kwargs"""
return {self.field_aliases.get(key, key): val for key, val in kwargs.items()}
[docs] def _unalias_kwargs_with_lookups(self, **kwargs):
"""convert alias name to solr field for keys in kwargs
with support for __ lookups for filters"""
new_kwargs = {}
for key, val in kwargs.items():
field_parts = key.split(self.LOOKUP_SEP, 1)
# first part is always present = field name
field = field_parts[0]
# get alias for key if there is one
field = self.field_aliases.get(field, field)
# if there is a lookup, add it back to the unaliased field
if len(field_parts) > 1:
field = "%s__%s" % (field, field_parts[1])
new_kwargs[field] = val
return new_kwargs
[docs] def filter(self, *args, tag: str = "", **kwargs) -> "AliasedSolrQuerySet":
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.filter`
to support using aliased field names for keyword argument keys."""
kwargs = self._unalias_kwargs_with_lookups(**kwargs)
return super().filter(*args, tag=tag, **kwargs)
[docs] def facet(self, *args, **kwargs) -> "AliasedSolrQuerySet":
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.facet`
to support using aliased field names in args."""
args = self._unalias_args(*args)
return super().facet(*args, **kwargs)
[docs] def stats(self, *args, **kwargs) -> "AliasedSolrQuerySet":
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.stats`
to support using aliased field names in args."""
args = self._unalias_args(*args)
return super().stats(*args, **kwargs)
[docs] def facet_field(
self, field: str, exclude: str = "", **kwargs
) -> "AliasedSolrQuerySet":
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.facet_field``
to support using aliased field names for field parameter."""
field = self.field_aliases.get(field, field)
return super().facet_field(field, exclude=exclude, **kwargs)
[docs] def order_by(self, *args) -> "AliasedSolrQuerySet":
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.order_by``
to support using aliased field names in sort arguments."""
args = self._unalias_args(*args)
return super().order_by(*args)
[docs] def only(self, *args, **kwargs) -> "AliasedSolrQuerySet":
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.only``
to support using aliased field names for args (but not kwargs)."""
# convert args to aliased args, switching them to keyword
# args; unknown fields are treated the same way
kwargs.update({arg: self.field_aliases.get(arg, arg) for arg in args})
return super().only(**kwargs)
# also method does not need to be extended, since it runs through only
[docs] def highlight(self, field: str, **kwargs) -> "AliasedSolrQuerySet":
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.highlight`
to support using aliased field names in kwargs."""
field = self.field_aliases.get(field, field)
return super().highlight(field, **kwargs)
[docs] def group(self, field: str, **kwargs) -> "AliasedSolrQuerySet":
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.group`
to support using aliased field names in kwargs. (Note that sorting
does not currently support aliased field names)."""
field = self.field_aliases.get(field, field)
# TODO: should we also reverse alias for sort option if specified?
return super().group(field, **kwargs)
[docs] def get_facets(self) -> Dict[str, int]:
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.get_facets`
to use aliased field names for facet and range facet keys."""
facets = super().get_facets()
# bail out if empty dict is returned
if not facets:
return facets
# replace field names in facet field and facet range
# with aliased field names
for section in ["facet_fields", "facet_ranges"]:
facets[section] = {
self.reverse_aliases.get(field, field): val
for field, val in facets[section].items()
}
return facets
[docs] def get_stats(self) -> Dict[str, Dict]:
"""Extend :meth:`parasolr.query.queryset.SolrQuerySet.get_stats` to
return return aliased field names for field_list keys."""
stats = super().get_stats()
if stats:
stats["stats_fields"] = {
self.reverse_aliases.get(field, field): val
for field, val in stats["stats_fields"].items()
}
return stats
[docs] def get_highlighting(self) -> Dict[str, Dict[str, List]]:
highlighting = super().get_highlighting()
# highlighting results are keyed on document id
# for each document, there is a dictionary of highlights;
# key is field name, value is the list of snippets
if highlighting:
for doc_id, highlights in highlighting.items():
highlighting[doc_id] = {
self.reverse_aliases.get(field, field): snippets
for field, snippets in highlights.items()
}
return highlighting
