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