test/test_Highlighter.py

   1 # ====================================================================
   2 #   Licensed under the Apache License, Version 2.0 (the "License");
   3 #   you may not use this file except in compliance with the License.
   4 #   You may obtain a copy of the License at
   5 #
   6 #       http://www.apache.org/licenses/LICENSE-2.0
   7 #
   8 #   Unless required by applicable law or agreed to in writing, software
   9 #   distributed under the License is distributed on an "AS IS" BASIS,
  10 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  11 #   See the License for the specific language governing permissions and
  12 #   limitations under the License.
  13 # ====================================================================
  14
  15 from unittest import TestCase, main
  16 from lucene import *
  17
  18
  19 class TestFormatter(PythonFormatter):
  20
  21     def __init__(self, testCase):
  22         super(TestFormatter, self).__init__()
  23         self.testCase = testCase
  24
  25     def highlightTerm(self, originalText, group):
  26         if group.getTotalScore() <= 0:
  27             return originalText;
  28
  29         self.testCase.countHighlightTerm()
  30
  31         return "<b>" + originalText + "</b>"
  32
  33
  34 class HighlighterTestCase(TestCase):
  35     """
  36     Unit tests ported from Java Lucene.
  37     2004 by Yura Smolsky ;)
  38     """
  39
  40     FIELD_NAME = "contents"
  41     texts = [ "A wicked problem is one for which each attempt to create a solution changes the understanding of the problem.  Wicked problems cannot be solved in a traditional linear fashion, because the problem definition evolves as new possible solutions are considered and/or implemented."
  42               "Wicked problems always occur in a social context -- the wickedness of the problem reflects the diversity among the stakeholders in the problem."
  43               "From http://cognexus.org/id42.htm"
  44               "Most projects in organizations -- and virtually all technology-related projects these days -- are about wicked problems.  Indeed, it is the social complexity of these problems, not their technical complexity, that overwhelms most current problem solving and project management approaches."
  45               "This text has a typo in referring to whicked problems" ];
  46
  47     def __init__(self, *args):
  48
  49         super(HighlighterTestCase, self).__init__(*args)
  50         self.parser = QueryParser(Version.LUCENE_CURRENT, self.FIELD_NAME,
  51                                   StandardAnalyzer(Version.LUCENE_CURRENT))
  52
  53     def testSimpleHighlighter(self):
  54
  55         self.doSearching("Wicked")
  56         highlighter = Highlighter(QueryScorer(self.query))
  57         highlighter.setTextFragmenter(SimpleFragmenter(40))
  58         maxNumFragmentsRequired = 2
  59
  60         for scoreDoc in self.scoreDocs:
  61             text = self.searcher.doc(scoreDoc.doc).get(self.FIELD_NAME)
  62             tokenStream = self.analyzer.tokenStream(self.FIELD_NAME,
  63                                                     StringReader(text))
  64
  65             result = highlighter.getBestFragments(tokenStream, text,
  66                                                   maxNumFragmentsRequired,
  67                                                   "...")
  68             print "\t", result
  69
  70         # Not sure we can assert anything here - just running to check we don't
  71         # throw any exceptions
  72
  73     def testGetBestFragmentsSimpleQuery(self):
  74
  75         self.doSearching("Wicked")
  76         self.doStandardHighlights()
  77         self.assert_(self.numHighlights == 3,
  78                      ("Failed to find correct number of highlights, %d found"
  79                       %(self.numHighlights)))
  80
  81     def doSearching(self, queryString):
  82
  83         self.searcher = IndexSearcher(self.ramDir, True)
  84         self.query = self.parser.parse(queryString)
  85         # for any multi-term queries to work (prefix, wildcard, range,
  86         # fuzzy etc) you must use a rewritten query!
  87         self.query = self.query.rewrite(self.reader)
  88
  89         print "Searching for:", self.query.toString(self.FIELD_NAME)
  90         self.scoreDocs = self.searcher.search(self.query, 100).scoreDocs
  91         self.numHighlights = 0
  92
  93     def doStandardHighlights(self):
  94
  95         formatter = TestFormatter(self)
  96
  97         highlighter = Highlighter(formatter, QueryScorer(self.query))
  98         highlighter.setTextFragmenter(SimpleFragmenter(20))
  99         for scoreDoc in self.scoreDocs:
 100             text = self.searcher.doc(scoreDoc.doc).get(self.FIELD_NAME)
 101             maxNumFragmentsRequired = 2
 102             fragmentSeparator = "..."
 103             tokenStream = self.analyzer.tokenStream(self.FIELD_NAME,
 104                                                     StringReader(text))
 105
 106             result = highlighter.getBestFragments(tokenStream,
 107                                                   text,
 108                                                   maxNumFragmentsRequired,
 109                                                   fragmentSeparator)
 110             print "\t", result
 111
 112     def countHighlightTerm(self):
 113
 114         self.numHighlights += 1 # update stats used in assertions
 115
 116     def setUp(self):
 117
 118         self.analyzer = StandardAnalyzer(Version.LUCENE_CURRENT)
 119         self.ramDir = RAMDirectory()
 120         writer = IndexWriter(self.ramDir, self.analyzer, True,
 121                              IndexWriter.MaxFieldLength.LIMITED)
 122         for text in self.texts:
 123             self.addDoc(writer, text)
 124
 125         writer.optimize()
 126         writer.close()
 127         self.reader = IndexReader.open(self.ramDir, True)
 128         self.numHighlights = 0;
 129
 130     def addDoc(self, writer, text):
 131
 132         d = Document()
 133         f = Field(self.FIELD_NAME, text,
 134                   Field.Store.YES, Field.Index.ANALYZED,
 135                   Field.TermVector.YES)
 136         d.add(f)
 137         writer.addDocument(d)
 138
 139
 140 if __name__ == "__main__":
 141     import sys, lucene
 142     lucene.initVM()
 143     if '-loop' in sys.argv:
 144         sys.argv.remove('-loop')
 145         while True:
 146             try:
 147                 main()
 148             except:
 149                 pass
 150     else:
 151          main()