apps/search/custom.py

   1
   2 from sunburnt import sunburnt
   3 from lxml import etree
   4 import urllib
   5 import warnings
   6 from sunburnt import search
   7 import copy
   8
   9
  10 class TermVectorOptions(search.Options):
  11     option_name = "tv"
  12
  13     def __init__(self, schema, original=None):
  14         self.schema = schema
  15         if original is None:
  16             self.fields = set()
  17             self.positions = False
  18         else:
  19             self.fields = copy.copy(original.fields)
  20             self.positions = copy.copy(original.positions)
  21
  22     def update(self, positions=False, fields=None):
  23         if fields is None:
  24             fields = []
  25         if isinstance(fields, basestring):
  26             fields = [fields]
  27         self.schema.check_fields(fields, {"stored": True})
  28         self.fields.update(fields)
  29         self.positions = positions
  30
  31     def options(self):
  32         opts = {}
  33         opts['tv'] = 'true'
  34         if self.positions:
  35             opts['tv.positions'] = 'true'
  36         if self.fields:
  37             opts['tv.fl'] = ','.join(sorted(self.fields))
  38         return opts
  39
  40
  41 class CustomSolrConnection(sunburnt.SolrConnection):
  42     def __init__(self, *args, **kw):
  43         super(CustomSolrConnection, self).__init__(*args, **kw)
  44         self.analysis_url = self.url + "analysis/field/"
  45
  46     def analyze(self, params):
  47         qs = urllib.urlencode(params)
  48         url = "%s?%s" % (self.analysis_url, qs)
  49         if len(url) > self.max_length_get_url:
  50             warnings.warn("Long query URL encountered - POSTing instead of "
  51                 "GETting. This query will not be cached at the HTTP layer")
  52             url = self.analysis_url
  53             kwargs = dict(
  54                 method="POST",
  55                 body=qs,
  56                 headers={"Content-Type": "application/x-www-form-urlencoded"},
  57             )
  58         else:
  59             kwargs = dict(method="GET")
  60         r, c = self.request(url, **kwargs)
  61         if r.status != 200:
  62             raise sunburnt.SolrError(r, c)
  63         return c
  64
  65
  66 # monkey patching sunburnt SolrSearch
  67 search.SolrSearch.option_modules += ('term_vectorer',)
  68
  69
  70 def __term_vector(self, positions=False, fields=None):
  71     newself = self.clone()
  72     newself.term_vectorer.update(positions, fields)
  73     return newself
  74 setattr(search.SolrSearch, 'term_vector', __term_vector)
  75 __original__init_common_modules = search.SolrSearch._init_common_modules
  76
  77
  78 def __patched__init_common_modules(self):
  79     __original__init_common_modules(self)
  80     self.term_vectorer = TermVectorOptions(self.schema)
  81 setattr(search.SolrSearch, '_init_common_modules', __patched__init_common_modules)
  82
  83
  84 class CustomSolrInterface(sunburnt.SolrInterface):
  85     # just copied from parent and SolrConnection -> CustomSolrConnection
  86     def __init__(self, url, schemadoc=None, http_connection=None, mode='', retry_timeout=-1, max_length_get_url=sunburnt.MAX_LENGTH_GET_URL):
  87         self.conn = CustomSolrConnection(url, http_connection, retry_timeout, max_length_get_url)
  88         self.schemadoc = schemadoc
  89         if mode == 'r':
  90             self.writeable = False
  91         elif mode == 'w':
  92             self.readable = False
  93         self.init_schema()
  94
  95     def _analyze(self, **kwargs):
  96         if not self.readable:
  97             raise TypeError("This Solr instance is only for writing")
  98         args = {
  99             'analysis_showmatch': True
 100             }
 101         if 'field' in kwargs: args['analysis_fieldname'] = kwargs['field']
 102         if 'text' in kwargs: args['analysis_fieldvalue'] = kwargs['text']
 103         if 'q' in kwargs: args['q'] = kwargs['q']
 104         if 'query' in kwargs: args['q'] = kwargs['q']
 105
 106         params = map(lambda (k, v): (k.replace('_', '.'), v), sunburnt.params_from_dict(**args))
 107
 108         content = self.conn.analyze(params)
 109         doc = etree.fromstring(content)
 110         return doc
 111
 112     def highlight(self, **kwargs):
 113         doc = self._analyze(**kwargs)
 114         analyzed = doc.xpath("//lst[@name='index']/arr[last()]/lst[bool/@name='match']")
 115         matches = set()
 116         for wrd in analyzed:
 117             start = int(wrd.xpath("int[@name='start']")[0].text)
 118             end = int(wrd.xpath("int[@name='end']")[0].text)
 119             matches.add((start, end))
 120
 121         print matches
 122         if matches:
 123             return self.substring(kwargs['text'], matches,
 124                             margins=kwargs.get('margins', 30),
 125             mark=kwargs.get('mark', ("<b>", "</b>")))
 126         else:
 127             return None
 128
 129     def analyze(self, **kwargs):
 130         doc = self._analyze(self, **kwargs)
 131         terms = doc.xpath("/lst[@name='index']/arr[last()]/lst/str[1]")
 132         terms = map(lambda n: unicode(n.text), terms)
 133         return terms
 134
 135     def substring(self, text, matches, margins=30, mark=("<b>", "</b>")):
 136         start = None
 137         end = None
 138         totlen = len(text)
 139         matches_margins = map(lambda (s, e): (max(0, s - margins), min(totlen, e + margins)), matches)
 140         (start, end) = matches_margins[0]
 141
 142         for (s, e) in matches_margins[1:]:
 143             if end < s or start > e:
 144                 continue
 145             start = min(start, s)
 146             end = max(end, e)
 147
 148         snip = text[start:end]
 149         matches = list(matches)
 150         matches.sort(lambda a, b: cmp(b[0], a[0]))
 151         for (s, e) in matches:
 152             off = - start
 153             snip = snip[:e + off] + mark[1] + snip[e + off:]
 154             snip = snip[:s + off] + mark[0] + snip[s + off:]
 155             # maybe break on word boundaries
 156         return snip