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