remove set_trace
[wolnelektury.git] / apps / search / index.py
index 9617077..47583ec 100644 (file)
@@ -27,6 +27,7 @@ from librarian import dcparser
 from librarian.parser import WLDocument
 from lxml import etree
 import catalogue.models
+from pdcounter.models import Author as PDCounterAuthor
 from multiprocessing.pool import ThreadPool
 from threading import current_thread
 import atexit
@@ -219,6 +220,15 @@ class Index(BaseIndex):
             doc.add(Field("tag_category", tag.category, Field.Store.NO, Field.Index.NOT_ANALYZED))
             self.index.addDocument(doc)
 
+        for pdtag in PDCounterAuthor.objects.all():
+            doc = Document()
+            doc.add(NumericField("tag_id", Field.Store.YES, True).setIntValue(int(pdtag.id)))
+            doc.add(Field("tag_name", pdtag.name, Field.Store.NO, Field.Index.ANALYZED))
+            doc.add(Field("tag_name_pl", pdtag.name, Field.Store.NO, Field.Index.ANALYZED))
+            doc.add(Field("tag_category", 'pdcounter', Field.Store.NO, Field.Index.NOT_ANALYZED))
+            doc.add(Field("is_pdcounter", 'true', Field.Store.YES, Field.Index.NOT_ANALYZED))
+            self.index.addDocument(doc)
+
     def create_book_doc(self, book):
         """
         Create a lucene document referring book id.
@@ -276,7 +286,7 @@ class Index(BaseIndex):
 
     footnote_tags = ['pa', 'pt', 'pr', 'pe']
 
-    skip_header_tags = ['autor_utworu', 'nazwa_utworu', 'dzielo_nadrzedne']
+    skip_header_tags = ['autor_utworu', 'nazwa_utworu', 'dzielo_nadrzedne', '{http://www.w3.org/1999/02/22-rdf-syntax-ns#}RDF']
 
     published_date_re = re.compile("([0-9]+)[\]. ]*$")
 
@@ -322,10 +332,10 @@ class Index(BaseIndex):
 
         # get published date
         source = book_info.source_name
-        match = self.published_date_re.search(source)
-        print("published date is %s %s" % (match, match is not None and match.groups()))
-        if match is not None:
-            fields["published_date"] = Field("published_date", str(match.groups()[0]), Field.Store.YES, Field.Index.NOT_ANALYZED)
+        if hasattr(book_info, 'source_name'):
+            match = self.published_date_re.search(source)
+            if match is not None:
+                fields["published_date"] = Field("published_date", str(match.groups()[0]), Field.Store.YES, Field.Index.NOT_ANALYZED)
 
         return fields
 
@@ -426,9 +436,8 @@ class Index(BaseIndex):
 
         fragments = {}
         snippets = Snippets(book.id).open('w')
-        position = 0
         try:
-            for header in list(master):
+            for header, position in zip(list(master), range(len(master))):
 
                 if header.tag in self.skip_header_tags:
                     continue
@@ -441,15 +450,15 @@ class Index(BaseIndex):
 
                 for start, end in walker(header, ignore_tags=self.ignore_content_tags):
                     # handle footnotes
-                    if start is not None and start.tag in self.footnote_tags:
-                        footnote = ' '.join(start.itertext())
-                    elif end is not None and footnote is not None and end.tag in self.footnote_tags:
-                        doc = add_part(snippets, header_index=position, header_type=header.tag,
-                                       content=footnote)
+                    if start is not None and start.tag in self.footnote_tags:
+                        footnote = ' '.join(start.itertext())
+                    elif end is not None and footnote is not None and end.tag in self.footnote_tags:
+                        doc = add_part(snippets, header_index=position, header_type=header.tag,
+                                       content=footnote)
 
-                        self.index.addDocument(doc)
+                        self.index.addDocument(doc)
 
-                        footnote = None
+                        footnote = None
 
                     # handle fragments and themes.
                     if start is not None and start.tag == 'begin':
@@ -496,7 +505,6 @@ class Index(BaseIndex):
                                content=fix_format(content))
 
                 self.index.addDocument(doc)
-                position += 1
 
         finally:
             snippets.close()
@@ -607,27 +615,25 @@ class SearchResult(object):
         stored = search.searcher.doc(scoreDocs.doc)
         self.book_id = int(stored.get("book_id"))
 
-        header_type = stored.get("header_type")
-        if not header_type:
-            return
-
-        sec = (header_type, int(stored.get("header_index")))
-        header_span = stored.get('header_span')
-        header_span = header_span is not None and int(header_span) or 1
-
-        fragment = stored.get("fragment_anchor")
-
         pd = stored.get("published_date")
         if pd is None:
-            print "published_date is none for book %d" % self.book_id
             pd = 0
         self.published_date = int(pd)
 
-        if snippets:
-            snippets = snippets.replace("/\n", "\n")
-        hit = (sec + (header_span,), fragment, scoreDocs.score, {'how_found': how_found, 'snippets': snippets and [snippets] or []})
+        header_type = stored.get("header_type")
+        # we have a content hit in some header of fragment
+        if header_type is not None:
+            sec = (header_type, int(stored.get("header_index")))
+            header_span = stored.get('header_span')
+            header_span = header_span is not None and int(header_span) or 1
+
+            fragment = stored.get("fragment_anchor")
 
-        self._hits.append(hit)
+            if snippets:
+                snippets = snippets.replace("/\n", "\n")
+            hit = (sec + (header_span,), fragment, scoreDocs.score, {'how_found': how_found, 'snippets': snippets and [snippets] or []})
+
+            self._hits.append(hit)
 
         self.search = search
         self.searched = searched
@@ -1212,19 +1218,27 @@ class Search(IndexStore):
         if terms:
             return JArray('object')(terms, Term)
 
-    def search_tags(self, query, filter=None, max_results=40):
+    def search_tags(self, query, filters=None, max_results=40, pdcounter=False):
         """
         Search for Tag objects using query.
         """
-        tops = self.searcher.search(query, filter, max_results)
+        if not pdcounter:
+            filters = self.chain_filters([filter, self.term_filter(Term('is_pdcounter', 'true'), inverse=True)])
+        tops = self.searcher.search(query, filters, max_results)
 
         tags = []
         for found in tops.scoreDocs:
             doc = self.searcher.doc(found.doc)
-            tag = catalogue.models.Tag.objects.get(id=doc.get("tag_id"))
-            tags.append(tag)
-            print "%s (%d) -> %f" % (tag, tag.id, found.score)
-
+            is_pdcounter = doc.get('is_pdcounter')
+            if is_pdcounter:
+                tag = PDCounterAuthor.objects.get(id=doc.get('tag_id'))
+            else:
+                tag = catalogue.models.Tag.objects.get(id=doc.get("tag_id"))
+                # don't add the pdcounter tag if same tag already exists
+            if not (is_pdcounter and filter(lambda t: tag.slug == t.slug, tags)):
+                tags.append(tag)
+                #            print "%s (%d) -> %f" % (tag, tag.id, found.score)
+        print 'returning %s' % tags
         return tags
 
     def search_books(self, query, filter=None, max_results=10):
@@ -1238,7 +1252,7 @@ class Search(IndexStore):
             bks.append(catalogue.models.Book.objects.get(id=doc.get("book_id")))
         return bks
 
-    def create_prefix_phrase(self, toks, field):
+    def make_prefix_phrase(self, toks, field):
         q = MultiPhraseQuery()
         for i in range(len(toks)):
             t = Term(field, toks[i])
@@ -1264,7 +1278,7 @@ class Search(IndexStore):
 
         return only_term
 
-    def hint_tags(self, string, max_results=50):
+    def hint_tags(self, string, max_results=50, pdcounter=True, prefix=True):
         """
         Return auto-complete hints for tags
         using prefix search.
@@ -1273,14 +1287,17 @@ class Search(IndexStore):
         top = BooleanQuery()
 
         for field in ['tag_name', 'tag_name_pl']:
-            q = self.create_prefix_phrase(toks, field)
+            if prefix:
+                q = self.make_prefix_phrase(toks, field)
+            else:
+                q = self.make_term_query(toks, field)
             top.add(BooleanClause(q, BooleanClause.Occur.SHOULD))
 
         no_book_cat = self.term_filter(Term("tag_category", "book"), inverse=True)
 
-        return self.search_tags(top, no_book_cat, max_results=max_results)
+        return self.search_tags(top, no_book_cat, max_results=max_results, pdcounter=pdcounter)
 
-    def hint_books(self, string, max_results=50):
+    def hint_books(self, string, max_results=50, prefix=True):
         """
         Returns auto-complete hints for book titles
         Because we do not index 'pseudo' title-tags.
@@ -1288,7 +1305,10 @@ class Search(IndexStore):
         """
         toks = self.get_tokens(string, field='SIMPLE')
 
-        q = self.create_prefix_phrase(toks, 'title')
+        if prefix:
+            q = self.make_prefix_phrase(toks, 'title')
+        else:
+            q = self.make_term_query(toks, 'title')
 
         return self.search_books(q, self.term_filter(Term("is_book", "true")), max_results=max_results)