apps/search/index.py

   1 # -*- coding: utf-8 -*-
   2
   3 from django.conf import settings
   4 from lucene import SimpleFSDirectory, IndexWriter, CheckIndex, \
   5     File, Field, Integer, \
   6     NumericField, Version, Document, JavaError, IndexSearcher, \
   7     QueryParser, PerFieldAnalyzerWrapper, \
   8     SimpleAnalyzer, PolishAnalyzer, ArrayList, \
   9     KeywordAnalyzer, NumericRangeQuery, NumericRangeFilter, BooleanQuery, \
  10     BlockJoinQuery, BlockJoinCollector, Filter, TermsFilter, ChainedFilter, \
  11     HashSet, BooleanClause, Term, CharTermAttribute, \
  12     PhraseQuery, MultiPhraseQuery, StringReader, TermQuery, \
  13     FuzzyQuery, FuzzyTermEnum, PrefixTermEnum, Sort, Integer, \
  14     SimpleHTMLFormatter, Highlighter, QueryScorer, TokenSources, TextFragment, \
  15     BooleanFilter, TermsFilter, FilterClause, QueryWrapperFilter, \
  16     initVM, CLASSPATH, JArray, JavaError
  17     # KeywordAnalyzer
  18
  19 # Initialize jvm
  20 JVM = initVM(CLASSPATH)
  21
  22 import sys
  23 import os
  24 import re
  25 import errno
  26 from librarian import dcparser
  27 from librarian.parser import WLDocument
  28 import catalogue.models
  29 from multiprocessing.pool import ThreadPool
  30 from threading import current_thread
  31 import atexit
  32 import traceback
  33
  34
  35 class WLAnalyzer(PerFieldAnalyzerWrapper):
  36     def __init__(self):
  37         polish = PolishAnalyzer(Version.LUCENE_34)
  38         #        polish_gap.setPositionIncrementGap(999)
  39
  40         simple = SimpleAnalyzer(Version.LUCENE_34)
  41         #        simple_gap.setPositionIncrementGap(999)
  42
  43         keyword = KeywordAnalyzer(Version.LUCENE_34)
  44
  45         # not sure if needed: there's NOT_ANALYZED meaning basically the same
  46
  47         PerFieldAnalyzerWrapper.__init__(self, polish)
  48
  49         self.addAnalyzer("tags", simple)
  50         self.addAnalyzer("technical_editors", simple)
  51         self.addAnalyzer("editors", simple)
  52         self.addAnalyzer("url", keyword)
  53         self.addAnalyzer("source_url", keyword)
  54         self.addAnalyzer("source_name", simple)
  55         self.addAnalyzer("publisher", simple)
  56         self.addAnalyzer("authors", simple)
  57         self.addAnalyzer("is_book", keyword)
  58         # shouldn't the title have two forms? _pl and simple?
  59
  60         self.addAnalyzer("themes", simple)
  61         self.addAnalyzer("themes_pl", polish)
  62
  63         self.addAnalyzer("tag_name", simple)
  64         self.addAnalyzer("tag_name_pl", polish)
  65
  66         self.addAnalyzer("translators", simple)
  67
  68         self.addAnalyzer("KEYWORD", keyword)
  69         self.addAnalyzer("SIMPLE", simple)
  70         self.addAnalyzer("POLISH", polish)
  71
  72
  73 class IndexStore(object):
  74     def __init__(self):
  75         self.make_index_dir()
  76         self.store = SimpleFSDirectory(File(settings.SEARCH_INDEX))
  77
  78     def make_index_dir(self):
  79         try:
  80             os.makedirs(settings.SEARCH_INDEX)
  81         except OSError as exc:
  82             if exc.errno == errno.EEXIST:
  83                 pass
  84             else: raise
  85
  86
  87 class IndexChecker(IndexStore):
  88     def __init__(self):
  89         IndexStore.__init__(self)
  90
  91     def check(self):
  92         checker = CheckIndex(self.store)
  93         status = checker.checkIndex()
  94         return status
  95
  96
  97 class Snippets(object):
  98     SNIPPET_DIR = "snippets"
  99
 100     def __init__(self, book_id):
 101         try:
 102             os.makedirs(os.path.join(settings.SEARCH_INDEX, self.SNIPPET_DIR))
 103         except OSError as exc:
 104             if exc.errno == errno.EEXIST:
 105                 pass
 106             else: raise
 107         self.book_id = book_id
 108         self.file = None
 109
 110     def open(self, mode='r'):
 111         if not 'b' in mode:
 112             mode += 'b'
 113         self.file = open(os.path.join(settings.SEARCH_INDEX, self.SNIPPET_DIR, str(self.book_id)), mode)
 114         self.position = 0
 115         return self
 116
 117     def add(self, snippet):
 118         txt = snippet.encode('utf-8')
 119         l = len(txt)
 120         self.file.write(txt)
 121         pos = (self.position, l)
 122         self.position += l
 123         return pos
 124
 125     def get(self, pos):
 126         self.file.seek(pos[0], 0)
 127         txt = self.file.read(pos[1]).decode('utf-8')
 128         return txt
 129
 130     def close(self):
 131         self.file.close()
 132
 133
 134 class Index(IndexStore):
 135     def __init__(self, analyzer=None):
 136         IndexStore.__init__(self)
 137         self.index = None
 138         if not analyzer:
 139             analyzer = WLAnalyzer()
 140         self.analyzer = analyzer
 141
 142     def open(self, analyzer=None):
 143         if self.index:
 144             raise Exception("Index is already opened")
 145         self.index = IndexWriter(self.store, self.analyzer,\
 146                                  IndexWriter.MaxFieldLength.LIMITED)
 147         return self.index
 148
 149     def optimize(self):
 150         self.index.optimize()
 151
 152     def close(self):
 153         try:
 154             self.index.optimize()
 155         except JavaError, je:
 156             print "Error during optimize phase, check index: %s" % je
 157
 158         self.index.close()
 159         self.index = None
 160
 161     def index_tags(self):
 162         q = NumericRangeQuery.newIntRange("tag_id", 0, Integer.MAX_VALUE, True, True)
 163         self.index.deleteDocuments(q)
 164
 165         for tag in catalogue.models.Tag.objects.all():
 166             doc = Document()
 167             doc.add(NumericField("tag_id", Field.Store.YES, True).setIntValue(tag.id))
 168             doc.add(Field("tag_name", tag.name, Field.Store.NO, Field.Index.ANALYZED))
 169             doc.add(Field("tag_name_pl", tag.name, Field.Store.NO, Field.Index.ANALYZED))
 170             doc.add(Field("tag_category", tag.category, Field.Store.NO, Field.Index.NOT_ANALYZED))
 171             self.index.addDocument(doc)
 172
 173     def remove_book(self, book):
 174         q = NumericRangeQuery.newIntRange("book_id", book.id, book.id, True, True)
 175         self.index.deleteDocuments(q)
 176
 177     def index_book(self, book, book_info=None, overwrite=True):
 178         if overwrite:
 179             self.remove_book(book)
 180
 181         book_doc = self.create_book_doc(book)
 182         meta_fields = self.extract_metadata(book, book_info)
 183         for f in meta_fields.values():
 184             if isinstance(f, list) or isinstance(f, tuple):
 185                 for elem in f:
 186                     book_doc.add(elem)
 187             else:
 188                 book_doc.add(f)
 189
 190         self.index.addDocument(book_doc)
 191         del book_doc
 192
 193         self.index_content(book, book_fields=[meta_fields['title'], meta_fields['authors']])
 194
 195     master_tags = [
 196         'opowiadanie',
 197         'powiesc',
 198         'dramat_wierszowany_l',
 199         'dramat_wierszowany_lp',
 200         'dramat_wspolczesny', 'liryka_l', 'liryka_lp',
 201         'wywiad'
 202         ]
 203
 204     skip_header_tags = ['autor_utworu', 'nazwa_utworu', 'dzielo_nadrzedne']
 205
 206     def create_book_doc(self, book):
 207         """
 208         Create a lucene document connected to the book
 209         """
 210         doc = Document()
 211         doc.add(NumericField("book_id", Field.Store.YES, True).setIntValue(book.id))
 212         if book.parent is not None:
 213             doc.add(NumericField("parent_id", Field.Store.YES, True).setIntValue(book.parent.id))
 214         return doc
 215
 216     def extract_metadata(self, book, book_info=None):
 217         fields = {}
 218
 219         if book_info is None:
 220             book_info = dcparser.parse(open(book.xml_file.path))
 221
 222         fields['slug'] = Field("slug", book.slug, Field.Store.NO, Field.Index.ANALYZED_NO_NORMS)
 223         fields['tags'] = self.add_gaps([Field("tags", t.name, Field.Store.NO, Field.Index.ANALYZED) for t in book.tags], 'tags')
 224         fields['is_book'] = Field("is_book", 'true', Field.Store.NO, Field.Index.NOT_ANALYZED)
 225
 226         # validator, name
 227         for field in dcparser.BookInfo.FIELDS:
 228             if hasattr(book_info, field.name):
 229                 if not getattr(book_info, field.name):
 230                     continue
 231                 # since no type information is available, we use validator
 232                 type_indicator = field.validator
 233                 if type_indicator == dcparser.as_unicode:
 234                     s = getattr(book_info, field.name)
 235                     if field.multiple:
 236                         s = ', '.join(s)
 237                     try:
 238                         fields[field.name] = Field(field.name, s, Field.Store.NO, Field.Index.ANALYZED)
 239                     except JavaError as je:
 240                         raise Exception("failed to add field: %s = '%s', %s(%s)" % (field.name, s, je.message, je.args))
 241                 elif type_indicator == dcparser.as_person:
 242                     p = getattr(book_info, field.name)
 243                     if isinstance(p, dcparser.Person):
 244                         persons = unicode(p)
 245                     else:
 246                         persons = ', '.join(map(unicode, p))
 247                     fields[field.name] = Field(field.name, persons, Field.Store.NO, Field.Index.ANALYZED)
 248                 elif type_indicator == dcparser.as_date:
 249                     dt = getattr(book_info, field.name)
 250                     fields[field.name] = Field(field.name, "%04d%02d%02d" %\
 251                                                (dt.year, dt.month, dt.day), Field.Store.NO, Field.Index.NOT_ANALYZED)
 252
 253         return fields
 254
 255     def get_master(self, root):
 256         for master in root.iter():
 257             if master.tag in self.master_tags:
 258                 return master
 259
 260     def add_gaps(self, fields, fieldname):
 261         def gap():
 262             while True:
 263                 yield Field(fieldname, ' ', Field.Store.NO, Field.Index.NOT_ANALYZED)
 264         return reduce(lambda a, b: a + b, zip(fields, gap()))[0:-1]
 265
 266     def index_content(self, book, book_fields=[]):
 267         wld = WLDocument.from_file(book.xml_file.path, parse_dublincore=False)
 268         root = wld.edoc.getroot()
 269
 270         master = self.get_master(root)
 271         if master is None:
 272             return []
 273
 274         def walker(node):
 275             yield node, None
 276             for child in list(node):
 277                 for b, e in walker(child):
 278                     yield b, e
 279             yield None, node
 280             return
 281
 282         def fix_format(text):
 283             return re.sub("(?m)/$", "", text)
 284
 285         def add_part(snippets, **fields):
 286             doc = self.create_book_doc(book)
 287             for f in book_fields:
 288                 doc.add(f)
 289
 290             doc.add(NumericField('header_index', Field.Store.YES, True).setIntValue(fields["header_index"]))
 291             doc.add(NumericField("header_span", Field.Store.YES, True)\
 292                     .setIntValue('header_span' in fields and fields['header_span'] or 1))
 293             doc.add(Field('header_type', fields["header_type"], Field.Store.YES, Field.Index.NOT_ANALYZED))
 294
 295             doc.add(Field('content', fields["content"], Field.Store.NO, Field.Index.ANALYZED, \
 296                           Field.TermVector.WITH_POSITIONS_OFFSETS))
 297
 298             snip_pos = snippets.add(fields["content"])
 299             doc.add(NumericField("snippets_position", Field.Store.YES, True).setIntValue(snip_pos[0]))
 300             doc.add(NumericField("snippets_length", Field.Store.YES, True).setIntValue(snip_pos[1]))
 301
 302             if 'fragment_anchor' in fields:
 303                 doc.add(Field("fragment_anchor", fields['fragment_anchor'],
 304                               Field.Store.YES, Field.Index.NOT_ANALYZED))
 305
 306             if 'themes' in fields:
 307                 themes, themes_pl = zip(*[
 308                     (Field("themes", theme, Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS),
 309                      Field("themes_pl", theme, Field.Store.NO, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS))
 310                      for theme in fields['themes']])
 311
 312                 themes = self.add_gaps(themes, 'themes')
 313                 themes_pl = self.add_gaps(themes_pl, 'themes_pl')
 314
 315                 for t in themes:
 316                     doc.add(t)
 317                 for t in themes_pl:
 318                     doc.add(t)
 319
 320             return doc
 321
 322         def give_me_utf8(s):
 323             if isinstance(s, unicode):
 324                 return s.encode('utf-8')
 325             else:
 326                 return s
 327
 328
 329         fragments = {}
 330         snippets = Snippets(book.id).open('w')
 331         try:
 332             for header, position in zip(list(master), range(len(master))):
 333
 334                 if header.tag in self.skip_header_tags:
 335                     continue
 336
 337                 content = u' '.join([t for t in header.itertext()])
 338                 content = fix_format(content)
 339
 340                 doc = add_part(snippets, header_index=position, header_type=header.tag, content=content)
 341
 342                 self.index.addDocument(doc)
 343
 344                 for start, end in walker(header):
 345                     if start is not None and start.tag == 'begin':
 346                         fid = start.attrib['id'][1:]
 347                         fragments[fid] = {'content': [], 'themes': [], 'start_section': position, 'start_header': header.tag}
 348                         fragments[fid]['content'].append(start.tail)
 349                     elif start is not None and start.tag == 'motyw':
 350                         fid = start.attrib['id'][1:]
 351                         if start.text is not None:
 352                             fragments[fid]['themes'] += map(str.strip, map(give_me_utf8, start.text.split(',')))
 353                         fragments[fid]['content'].append(start.tail)
 354                     elif start is not None and start.tag == 'end':
 355                         fid = start.attrib['id'][1:]
 356                         if fid not in fragments:
 357                             continue  # a broken <end> node, skip it
 358                         frag = fragments[fid]
 359                         if frag['themes'] == []:
 360                             continue  # empty themes list.
 361                         del fragments[fid]
 362
 363                         def jstr(l):
 364                             return u' '.join(map(
 365                                 lambda x: x == None and u'(none)' or unicode(x),
 366                                 l))
 367
 368                         doc = add_part(snippets,
 369                                        header_type=frag['start_header'],
 370                                        header_index=frag['start_section'],
 371                                        header_span=position - frag['start_section'] + 1,
 372                                        fragment_anchor=fid,
 373                                        content=u' '.join(filter(lambda s: s is not None, frag['content'])),
 374                                        themes=frag['themes'])
 375
 376                         self.index.addDocument(doc)
 377                     elif start is not None:
 378                         for frag in fragments.values():
 379                             frag['content'].append(start.text)
 380                     elif end is not None:
 381                         for frag in fragments.values():
 382                             frag['content'].append(end.tail)
 383         finally:
 384             snippets.close()
 385
 386
 387     def __enter__(self):
 388         self.open()
 389         return self
 390
 391     def __exit__(self, type, value, tb):
 392         self.close()
 393
 394
 395 def log_exception_wrapper(f):
 396     def _wrap(*a):
 397         try:
 398             f(*a)
 399         except Exception, e:
 400             print("Error in indexing thread: %s" % e)
 401             traceback.print_exc()
 402             raise e
 403     return _wrap
 404
 405
 406 class ReusableIndex(Index):
 407     """
 408     Works like index, but does not close/optimize Lucene index
 409     until program exit (uses atexit hook).
 410     This is usefull for importbooks command.
 411
 412     if you cannot rely on atexit, use ReusableIndex.close_reusable() yourself.
 413     """
 414     index = None
 415
 416     def open(self, analyzer=None, threads=4):
 417         if ReusableIndex.index is not None:
 418             self.index = ReusableIndex.index
 419         else:
 420             print("opening index")
 421             Index.open(self, analyzer)
 422             ReusableIndex.index = self.index
 423             atexit.register(ReusableIndex.close_reusable)
 424
 425     # def index_book(self, *args, **kw):
 426     #     job = ReusableIndex.pool.apply_async(log_exception_wrapper(Index.index_book), (self,) + args, kw)
 427     #     ReusableIndex.pool_jobs.append(job)
 428
 429     @staticmethod
 430     def close_reusable():
 431         if ReusableIndex.index is not None:
 432             ReusableIndex.index.optimize()
 433             ReusableIndex.index.close()
 434             ReusableIndex.index = None
 435
 436     def close(self):
 437         pass
 438
 439
 440 class Search(IndexStore):
 441     def __init__(self, default_field="content"):
 442         IndexStore.__init__(self)
 443         self.analyzer = WLAnalyzer() #PolishAnalyzer(Version.LUCENE_34)
 444         ## self.analyzer = WLAnalyzer()
 445         self.searcher = IndexSearcher(self.store, True)
 446         self.parser = QueryParser(Version.LUCENE_34, default_field,
 447                                   self.analyzer)
 448
 449         self.parent_filter = TermsFilter()
 450         self.parent_filter.addTerm(Term("is_book", "true"))
 451
 452     def query(self, query):
 453         return self.parser.parse(query)
 454
 455     def wrapjoins(self, query, fields=[]):
 456         """
 457         This functions modifies the query in a recursive way,
 458         so Term and Phrase Queries contained, which match
 459         provided fields are wrapped in a BlockJoinQuery,
 460         and so delegated to children documents.
 461         """
 462         if BooleanQuery.instance_(query):
 463             qs = BooleanQuery.cast_(query)
 464             for clause in qs:
 465                 clause = BooleanClause.cast_(clause)
 466                 clause.setQuery(self.wrapjoins(clause.getQuery(), fields))
 467             return qs
 468         else:
 469             termset = HashSet()
 470             query.extractTerms(termset)
 471             for t in termset:
 472                 t = Term.cast_(t)
 473                 if t.field() not in fields:
 474                     return query
 475             return BlockJoinQuery(query, self.parent_filter,
 476                                   BlockJoinQuery.ScoreMode.Total)
 477
 478     def simple_search(self, query, max_results=50):
 479         """Returns (books, total_hits)
 480         """
 481
 482         tops = self.searcher.search(self.query(query), max_results)
 483         bks = []
 484         for found in tops.scoreDocs:
 485             doc = self.searcher.doc(found.doc)
 486             bks.append(catalogue.models.Book.objects.get(id=doc.get("book_id")))
 487         return (bks, tops.totalHits)
 488
 489
 490     def search(self, query, max_results=50):
 491         query = self.query(query)
 492         query = self.wrapjoins(query, ["content", "themes"])
 493
 494         tops = self.searcher.search(query, max_results)
 495         bks = []
 496         for found in tops.scoreDocs:
 497             doc = self.searcher.doc(found.doc)
 498             bks.append(catalogue.models.Book.objects.get(id=doc.get("book_id")))
 499         return (bks, tops.totalHits)
 500
 501     def bsearch(self, query, max_results=50):
 502         q = self.query(query)
 503         bjq = BlockJoinQuery(q, self.parent_filter, BlockJoinQuery.ScoreMode.Avg)
 504
 505         tops = self.searcher.search(bjq, max_results)
 506         bks = []
 507         for found in tops.scoreDocs:
 508             doc = self.searcher.doc(found.doc)
 509             bks.append(catalogue.models.Book.objects.get(id=doc.get("book_id")))
 510         return (bks, tops.totalHits)
 511
 512 # TokenStream tokenStream = analyzer.tokenStream(fieldName, reader);
 513 # OffsetAttribute offsetAttribute = tokenStream.getAttribute(OffsetAttribute.class);
 514 # CharTermAttribute charTermAttribute = tokenStream.getAttribute(CharTermAttribute.class);
 515
 516 # while (tokenStream.incrementToken()) {
 517 #     int startOffset = offsetAttribute.startOffset();
 518 #     int endOffset = offsetAttribute.endOffset();
 519 #     String term = charTermAttribute.toString();
 520 # }
 521
 522
 523 class SearchResult(object):
 524     def __init__(self, searcher, scoreDocs, score=None, how_found=None, snippets=None):
 525         self.snippets = []
 526
 527         if score:
 528             self.score = score
 529         else:
 530             self.score = scoreDocs.score
 531
 532         self.hits = []
 533
 534         stored = searcher.doc(scoreDocs.doc)
 535         self.book_id = int(stored.get("book_id"))
 536
 537         header_type = stored.get("header_type")
 538         if not header_type:
 539             return
 540
 541         sec = (header_type, int(stored.get("header_index")))
 542         header_span = stored.get('header_span')
 543         header_span = header_span is not None and int(header_span) or 1
 544
 545         fragment = stored.get("fragment_anchor")
 546
 547         hit = (sec + (header_span,), fragment, scoreDocs.score, {'how_found': how_found, 'snippets': snippets})
 548
 549         self.hits.append(hit)
 550
 551     def merge(self, other):
 552         if self.book_id != other.book_id:
 553             raise ValueError("this search result is or book %d; tried to merge with %d" % (self.book_id, other.book_id))
 554         self.hits += other.hits
 555         if other.score > self.score:
 556             self.score = other.score
 557         return self
 558
 559     def get_book(self):
 560         return catalogue.models.Book.objects.get(id=self.book_id)
 561
 562     book = property(get_book)
 563
 564     def process_hits(self):
 565         frags = filter(lambda r: r[1] is not None, self.hits)
 566         sect = filter(lambda r: r[1] is None, self.hits)
 567         sect = filter(lambda s: 0 == len(filter(
 568             lambda f: s[0][1] >= f[0][1] and s[0][1] < f[0][1] + f[0][2],
 569             frags)), sect)
 570
 571         hits = []
 572
 573         for s in sect:
 574             m = {'score': s[2],
 575                  'header_index': s[0][1]
 576                  }
 577             m.update(s[3])
 578             hits.append(m)
 579
 580         for f in frags:
 581             frag = catalogue.models.Fragment.objects.get(anchor=f[1])
 582             m = {'score': f[2],
 583                  'fragment': frag,
 584                  'themes': frag.tags.filter(category='theme')
 585                  }
 586             m.update(f[3])
 587             hits.append(m)
 588
 589         hits.sort(lambda a, b: cmp(a['score'], b['score']), reverse=True)
 590
 591         print("--- %s" % hits)
 592
 593         return hits
 594
 595     def __unicode__(self):
 596         return u'SearchResult(book_id=%d, score=%d)' % (self.book_id, self.score)
 597
 598     @staticmethod
 599     def aggregate(*result_lists):
 600         books = {}
 601         for rl in result_lists:
 602             for r in rl:
 603                 if r.book_id in books:
 604                     books[r.book_id].merge(r)
 605                     #print(u"already have one with score %f, and this one has score %f" % (books[book.id][0], found.score))
 606                 else:
 607                     books[r.book_id] = r
 608         return books.values()
 609
 610     def __cmp__(self, other):
 611         return cmp(self.score, other.score)
 612
 613
 614 class Hint(object):
 615     def __init__(self, search):
 616         self.search = search
 617         self.book_tags = {}
 618         self.part_tags = []
 619         self._books = []
 620
 621     def books(self, *books):
 622         self._books = books
 623
 624     def tags(self, tags):
 625         for t in tags:
 626             if t.category in ['author', 'title', 'epoch', 'genre', 'kind']:
 627                 lst = self.book_tags.get(t.category, [])
 628                 lst.append(t)
 629                 self.book_tags[t.category] = lst
 630             if t.category in ['theme']:
 631                 self.part_tags.append(t)
 632
 633     def tag_filter(self, tags, field='tags'):
 634         q = BooleanQuery()
 635
 636         for tag in tags:
 637             toks = self.search.get_tokens(tag.name, field=field)
 638             tag_phrase = PhraseQuery()
 639             for tok in toks:
 640                 tag_phrase.add(Term(field, tok))
 641             q.add(BooleanClause(tag_phrase, BooleanClause.Occur.MUST))
 642
 643         return QueryWrapperFilter(q)
 644
 645     def book_filter(self):
 646         tags = reduce(lambda a, b: a + b, self.book_tags.values(), [])
 647         if tags:
 648             return self.tag_filter(tags)
 649         else:
 650             return None
 651
 652     def part_filter(self):
 653         fs = []
 654         if self.part_tags:
 655             fs.append(self.tag_filter(self.part_tags, field='themes'))
 656
 657         if self._books != []:
 658             bf = BooleanFilter()
 659             for b in self._books:
 660                 id_filter = NumericRangeFilter.newIntRange('book_id', b.id, b.id, True, True)
 661                 bf.add(FilterClause(id_filter, BooleanClause.Occur.SHOULD))
 662             fs.append(bf)
 663
 664         return MultiSearch.chain_filters(fs)
 665
 666     def should_search_for_book(self):
 667         return self._books == []
 668
 669     def just_search_in(self, all):
 670         """Holds logic to figure out which indexes should be search, when we have some hinst already"""
 671         some = []
 672         for field in all:
 673             if field == 'authors' and 'author' in self.book_tags:
 674                 continue
 675             if field == 'title' and self._books != []:
 676                 continue
 677             if (field == 'themes' or field == 'themes_pl') and self.part_tags:
 678                 continue
 679             some.append(field)
 680         return some
 681
 682
 683 class MultiSearch(Search):
 684     """Class capable of IMDb-like searching"""
 685     def get_tokens(self, searched, field='content'):
 686         """returns tokens analyzed by a proper (for a field) analyzer
 687         argument can be: StringReader, string/unicode, or tokens. In the last case
 688         they will just be returned (so we can reuse tokens, if we don't change the analyzer)
 689         """
 690         if isinstance(searched, str) or isinstance(searched, unicode):
 691             searched = StringReader(searched)
 692         elif isinstance(searched, list):
 693             return searched
 694
 695         searched.reset()
 696         tokens = self.analyzer.reusableTokenStream(field, searched)
 697         toks = []
 698         while tokens.incrementToken():
 699             cta = tokens.getAttribute(CharTermAttribute.class_)
 700             toks.append(cta.toString())
 701         return toks
 702
 703     def fuzziness(self, fuzzy):
 704         if not fuzzy:
 705             return None
 706         if isinstance(fuzzy, float) and fuzzy > 0.0 and fuzzy <= 1.0:
 707             return fuzzy
 708         else:
 709             return 0.5
 710
 711     def make_phrase(self, tokens, field='content', slop=2, fuzzy=False):
 712         if fuzzy:
 713             phrase = MultiPhraseQuery()
 714             for t in tokens:
 715                 term = Term(field, t)
 716                 fuzzterm = FuzzyTermEnum(self.searcher.getIndexReader(), term, self.fuzziness(fuzzy))
 717                 fuzzterms = []
 718
 719                 while True:
 720                     #                    print("fuzz %s" % unicode(fuzzterm.term()).encode('utf-8'))
 721                     ft = fuzzterm.term()
 722                     if ft:
 723                         fuzzterms.append(ft)
 724                     if not fuzzterm.next(): break
 725                 if fuzzterms:
 726                     phrase.add(JArray('object')(fuzzterms, Term))
 727                 else:
 728                     phrase.add(term)
 729         else:
 730             phrase = PhraseQuery()
 731             phrase.setSlop(slop)
 732             for t in tokens:
 733                 term = Term(field, t)
 734                 phrase.add(term)
 735         return phrase
 736
 737     def make_term_query(self, tokens, field='content', modal=BooleanClause.Occur.SHOULD, fuzzy=False):
 738         q = BooleanQuery()
 739         for t in tokens:
 740             term = Term(field, t)
 741             if fuzzy:
 742                 term = FuzzyQuery(term, self.fuzziness(fuzzy))
 743             else:
 744                 term = TermQuery(term)
 745             q.add(BooleanClause(term, modal))
 746         return q
 747
 748     # def content_query(self, query):
 749     #     return BlockJoinQuery(query, self.parent_filter,
 750     #                           BlockJoinQuery.ScoreMode.Total)
 751
 752     def search_perfect_book(self, searched, max_results=20, fuzzy=False, hint=None):
 753         fields_to_search = ['authors', 'title']
 754         only_in = None
 755         if hint:
 756             if not hint.should_search_for_book():
 757                 return []
 758             fields_to_search = hint.just_search_in(fields_to_search)
 759             only_in = hint.book_filter()
 760
 761         qrys = [self.make_phrase(self.get_tokens(searched, field=fld), field=fld, fuzzy=fuzzy) for fld in fields_to_search]
 762
 763         books = []
 764         for q in qrys:
 765             top = self.searcher.search(q,
 766                                        self.chain_filters([only_in, self.term_filter(Term('is_book', 'true'))]),
 767                 max_results)
 768             for found in top.scoreDocs:
 769                 books.append(SearchResult(self.searcher, found))
 770         return books
 771
 772     def search_perfect_parts(self, searched, max_results=20, fuzzy=False, hint=None):
 773         qrys = [self.make_phrase(self.get_tokens(searched), field=fld, fuzzy=fuzzy) for fld in ['content']]
 774
 775         flt = None
 776         if hint:
 777             flt = hint.part_filter()
 778
 779         books = []
 780         for q in qrys:
 781             top = self.searcher.search(q,
 782                                        self.chain_filters([self.term_filter(Term('is_book', 'true'), inverse=True),
 783                                                            flt
 784                                                           ]),
 785                                        max_results)
 786             for found in top.scoreDocs:
 787                 books.append(SearchResult(self.searcher, found, snippets=self.get_snippets(found, q)))
 788
 789         return books
 790
 791     def search_everywhere(self, searched, max_results=20, fuzzy=False, hint=None):
 792         books = []
 793         only_in = None
 794
 795         if hint:
 796             only_in = hint.part_filter()
 797
 798         # content only query : themes x content
 799         q = BooleanQuery()
 800
 801         tokens = self.get_tokens(searched)
 802         if hint is None or hint.just_search_in(['themes_pl']) != []:
 803             q.add(BooleanClause(self.make_term_query(tokens, field='themes_pl',
 804                                                      fuzzy=fuzzy), BooleanClause.Occur.MUST))
 805
 806         q.add(BooleanClause(self.make_term_query(tokens, field='content',
 807                                                  fuzzy=fuzzy), BooleanClause.Occur.SHOULD))
 808
 809         topDocs = self.searcher.search(q, only_in, max_results)
 810         for found in topDocs.scoreDocs:
 811             books.append(SearchResult(self.searcher, found))
 812
 813         # query themes/content x author/title/tags
 814         q = BooleanQuery()
 815         #        in_meta = BooleanQuery()
 816         in_content = BooleanQuery()
 817
 818         for fld in ['themes', 'content', 'tags', 'authors', 'title']:
 819             in_content.add(BooleanClause(self.make_term_query(tokens, field=fld, fuzzy=False), BooleanClause.Occur.SHOULD))
 820
 821         topDocs = self.searcher.search(q, only_in, max_results)
 822         for found in topDocs.scoreDocs:
 823             books.append(SearchResult(self.searcher, found))
 824
 825         return books
 826
 827     def multisearch(self, query, max_results=50):
 828         """
 829         Search strategy:
 830         - (phrase) OR -> content
 831                       -> title
 832                       -> authors
 833         - (keywords)  -> authors
 834                       -> motyw
 835                       -> tags
 836                       -> content
 837         """
 838         # queryreader = StringReader(query)
 839         # tokens = self.get_tokens(queryreader)
 840
 841         # top_level = BooleanQuery()
 842         # Should = BooleanClause.Occur.SHOULD
 843
 844         # phrase_level = BooleanQuery()
 845         # phrase_level.setBoost(1.3)
 846
 847         # p_content = self.make_phrase(tokens, joined=True)
 848         # p_title = self.make_phrase(tokens, 'title')
 849         # p_author = self.make_phrase(tokens, 'author')
 850
 851         # phrase_level.add(BooleanClause(p_content, Should))
 852         # phrase_level.add(BooleanClause(p_title, Should))
 853         # phrase_level.add(BooleanClause(p_author, Should))
 854
 855         # kw_level = BooleanQuery()
 856
 857         # kw_level.add(self.make_term_query(tokens, 'author'), Should)
 858         # j_themes = self.make_term_query(tokens, 'themes', joined=True)
 859         # kw_level.add(j_themes, Should)
 860         # kw_level.add(self.make_term_query(tokens, 'tags'), Should)
 861         # j_con = self.make_term_query(tokens, joined=True)
 862         # kw_level.add(j_con, Should)
 863
 864         # top_level.add(BooleanClause(phrase_level, Should))
 865         # top_level.add(BooleanClause(kw_level, Should))
 866
 867         return None
 868
 869     def book_search(self, query, filter=None, max_results=50, collector=None):
 870         tops = self.searcher.search(query, filter, max_results)
 871         #tops = self.searcher.search(p_content, max_results)
 872
 873         bks = []
 874         for found in tops.scoreDocs:
 875             doc = self.searcher.doc(found.doc)
 876             b = catalogue.models.Book.objects.get(id=doc.get("book_id"))
 877             bks.append(b)
 878             print "%s (%d) -> %f" % (b, b.id, found.score)
 879         return bks
 880
 881     def get_snippets(self, scoreDoc, query, field='content'):
 882         htmlFormatter = SimpleHTMLFormatter()
 883         highlighter = Highlighter(htmlFormatter, QueryScorer(query))
 884
 885         stored = self.searcher.doc(scoreDoc.doc)
 886
 887         # locate content.
 888         snippets = Snippets(stored.get('book_id')).open()
 889         try:
 890             text = snippets.get((int(stored.get('snippets_position')),
 891                                  int(stored.get('snippets_length'))))
 892         finally:
 893             snippets.close()
 894
 895         tokenStream = TokenSources.getAnyTokenStream(self.searcher.getIndexReader(), scoreDoc.doc, field, self.analyzer)
 896         #  highlighter.getBestTextFragments(tokenStream, text, False, 10)
 897         #        import pdb; pdb.set_trace()
 898         snip = highlighter.getBestFragments(tokenStream, text, 3, "...")
 899
 900         return [snip]
 901
 902     @staticmethod
 903     def enum_to_array(enum):
 904         """
 905         Converts a lucene TermEnum to array of Terms, suitable for
 906         addition to queries
 907         """
 908         terms = []
 909
 910         while True:
 911             t = enum.term()
 912             if t:
 913                 terms.append(t)
 914             if not enum.next(): break
 915
 916         if terms:
 917             return JArray('object')(terms, Term)
 918
 919     def search_tags(self, query, filter=None, max_results=40):
 920         tops = self.searcher.search(query, filter, max_results)
 921
 922         tags = []
 923         for found in tops.scoreDocs:
 924             doc = self.searcher.doc(found.doc)
 925             tag = catalogue.models.Tag.objects.get(id=doc.get("tag_id"))
 926             tags.append(tag)
 927             print "%s (%d) -> %f" % (tag, tag.id, found.score)
 928
 929         return tags
 930
 931     def search_books(self, query, filter=None, max_results=10):
 932         bks = []
 933         tops = self.searcher.search(query, filter, max_results)
 934         for found in tops.scoreDocs:
 935             doc = self.searcher.doc(found.doc)
 936             bks.append(catalogue.models.Book.objects.get(id=doc.get("book_id")))
 937         return bks
 938
 939     def create_prefix_phrase(self, toks, field):
 940         q = MultiPhraseQuery()
 941         for i in range(len(toks)):
 942             t = Term(field, toks[i])
 943             if i == len(toks) - 1:
 944                 pterms = MultiSearch.enum_to_array(PrefixTermEnum(self.searcher.getIndexReader(), t))
 945                 if pterms:
 946                     q.add(pterms)
 947                 else:
 948                     q.add(t)
 949             else:
 950                 q.add(t)
 951         return q
 952
 953     @staticmethod
 954     def term_filter(term, inverse=False):
 955         only_term = TermsFilter()
 956         only_term.addTerm(term)
 957
 958         if inverse:
 959             neg = BooleanFilter()
 960             neg.add(FilterClause(only_term, BooleanClause.Occur.MUST_NOT))
 961             only_term = neg
 962
 963         return only_term
 964
 965     def hint_tags(self, string, max_results=50):
 966         toks = self.get_tokens(string, field='SIMPLE')
 967         top = BooleanQuery()
 968
 969         for field in ['tag_name', 'tag_name_pl']:
 970             q = self.create_prefix_phrase(toks, field)
 971             top.add(BooleanClause(q, BooleanClause.Occur.SHOULD))
 972
 973         no_book_cat = self.term_filter(Term("tag_category", "book"), inverse=True)
 974
 975         return self.search_tags(top, no_book_cat, max_results=max_results)
 976
 977     def hint_books(self, string, max_results=50):
 978         toks = self.get_tokens(string, field='SIMPLE')
 979
 980         q = self.create_prefix_phrase(toks, 'title')
 981
 982         return self.book_search(q, self.term_filter(Term("is_book", "true")), max_results=max_results)
 983
 984     @staticmethod
 985     def chain_filters(filters, op=ChainedFilter.AND):
 986         filters = filter(lambda x: x is not None, filters)
 987         if not filters:
 988             return None
 989         chf = ChainedFilter(JArray('object')(filters, Filter), op)
 990         return chf
 991
 992     def filtered_categories(self, tags):
 993         cats = {}
 994         for t in tags:
 995             cats[t.category] = True
 996         return cats.keys()
 997
 998     def hint(self):
 999         return Hint(self)