samples/LuceneInAction/lia/analysis/synonym/SynonymFilter.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 lucene import Token, PythonTokenFilter, TermAttribute
  16 from lia.analysis.AnalyzerUtils import AnalyzerUtils
  17
  18 #
  19 # A TokenFilter extension
  20 #
  21
  22 class SynonymFilter(PythonTokenFilter):
  23     TOKEN_TYPE_SYNONYM = "SYNONYM"
  24
  25     def __init__(self, inStream, engine):
  26         super(SynonymFilter, self).__init__(inStream)
  27
  28         self.synonymStack = []
  29         self.termAttr = self.addAttribute(TermAttribute.class_)
  30         self.save = inStream.cloneAttributes()
  31         self.engine = engine
  32         self.inStream = inStream
  33
  34     def incrementToken(self):
  35
  36         if len(self.synonymStack) > 0:
  37             syn = self.synonymStack.pop()
  38             self.restoreState(syn)
  39             return True
  40
  41         if not self.inStream.incrementToken():
  42             return False
  43
  44         self.addAliasesToStack()
  45
  46         return True
  47
  48     def addAliasesToStack(self):
  49
  50         synonyms = self.engine.getSynonyms(self.termAttr.term())
  51         if synonyms is None:
  52             return
  53
  54         current = self.captureState()
  55
  56         for synonym in synonyms:
  57             self.save.restoreState(current)
  58             AnalyzerUtils.setTerm(self.save, synonym)
  59             AnalyzerUtils.setType(self.save, self.TOKEN_TYPE_SYNONYM)
  60             AnalyzerUtils.setPositionIncrement(self.save, 0)
  61             self.synonymStack.append(self.save.captureState())