cola/widgets/spellcheck.py

   1 #!/usr/bin/env python
   2 # -*- coding: utf-8 -*-
   3
   4 __copyright__ = """
   5 2012, Peter Norvig (http://norvig.com/spell-correct.html)
   6 2013, David Aguilar <davvid@gmail.com>
   7 """
   8
   9 import collections
  10 import re
  11 import sys
  12
  13 from PyQt4.Qt import QAction
  14 from PyQt4.Qt import QApplication
  15 from PyQt4.Qt import QEvent
  16 from PyQt4.Qt import QMenu
  17 from PyQt4.Qt import QMouseEvent
  18 from PyQt4.Qt import QSyntaxHighlighter
  19 from PyQt4.Qt import QTextCharFormat
  20 from PyQt4.Qt import QTextCursor
  21 from PyQt4.Qt import Qt
  22 from PyQt4.QtCore import SIGNAL
  23
  24 from cola.compat import set
  25 from cola.i18n import N_
  26 from cola.widgets.text import HintedTextEdit
  27
  28
  29 alphabet = 'abcdefghijklmnopqrstuvwxyz'
  30
  31
  32 def train(features, model):
  33     for f in features:
  34         model[f] += 1
  35     return model
  36
  37
  38 def edits1(word):
  39    splits     = [(word[:i], word[i:]) for i in range(len(word) + 1)]
  40    deletes    = [a + b[1:] for a, b in splits if b]
  41    transposes = [a + b[1] + b[0] + b[2:] for a, b in splits if len(b)>1]
  42    replaces   = [a + c + b[1:] for a, b in splits for c in alphabet if b]
  43    inserts    = [a + c + b     for a, b in splits for c in alphabet]
  44    return set(deletes + transposes + replaces + inserts)
  45
  46
  47 def known_edits2(word, words):
  48     return set(e2 for e1 in edits1(word)
  49                   for e2 in edits1(e1) if e2 in words)
  50
  51 def known(word, words):
  52     return set(w for w in word if w in words)
  53
  54
  55 def suggest(word, words):
  56     candidates = (known([word], words) or
  57                   known(edits1(word), words) or
  58                   known_edits2(word, words) or [word])
  59     return candidates
  60
  61
  62 def correct(word, words):
  63     candidates = suggest(word, words)
  64     return max(candidates, key=words.get)
  65
  66
  67 class NorvigSpellCheck(object):
  68     def __init__(self):
  69         self.words = collections.defaultdict(lambda: 1)
  70         self.extra_words = set()
  71         self.initialized = False
  72
  73     def init(self):
  74         if self.initialized:
  75             return
  76         self.initialized = True
  77         train(self.read(), self.words)
  78         train(self.extra_words, self.words)
  79
  80     def add_word(self, word):
  81         self.extra_words.add(word)
  82
  83     def suggest(self, word):
  84         self.init()
  85         return suggest(word, self.words)
  86
  87     def check(self, word):
  88         self.init()
  89         return word.replace('.', '') in self.words
  90
  91     def read(self):
  92         for (path, title) in (('/usr/share/dict/words', True),
  93                               ('/usr/share/dict/propernames', False)):
  94             try:
  95                 with open(path, 'r') as f:
  96                     for word in f:
  97                         yield word.rstrip()
  98                         if title:
  99                             yield word.rstrip().title()
 100             except IOError:
 101                 pass
 102         raise StopIteration
 103
 104
 105 class SpellCheckTextEdit(HintedTextEdit):
 106
 107     def __init__(self, hint, parent=None):
 108         HintedTextEdit.__init__(self, hint, parent)
 109
 110         # Default dictionary based on the current locale.
 111         self.spellcheck = NorvigSpellCheck()
 112         self.highlighter = Highlighter(self.document(), self.spellcheck)
 113
 114     def mousePressEvent(self, event):
 115         if event.button() == Qt.RightButton:
 116             # Rewrite the mouse event to a left button event so the cursor is
 117             # moved to the location of the pointer.
 118             event = QMouseEvent(QEvent.MouseButtonPress,
 119                                 event.pos(),
 120                                 Qt.LeftButton,
 121                                 Qt.LeftButton,
 122                                 Qt.NoModifier)
 123         HintedTextEdit.mousePressEvent(self, event)
 124
 125     def context_menu(self):
 126         popup_menu = HintedTextEdit.createStandardContextMenu(self)
 127
 128         # Select the word under the cursor.
 129         cursor = self.textCursor()
 130         cursor.select(QTextCursor.WordUnderCursor)
 131         self.setTextCursor(cursor)
 132
 133         # Check if the selected word is misspelled and offer spelling
 134         # suggestions if it is.
 135         spell_menu = None
 136         if self.textCursor().hasSelection():
 137             text = unicode(self.textCursor().selectedText())
 138             if not self.spellcheck.check(text):
 139                 spell_menu = QMenu(N_('Spelling Suggestions'))
 140                 for word in self.spellcheck.suggest(text):
 141                     action = SpellAction(word, spell_menu)
 142                     self.connect(action, SIGNAL('correct'), self.correct)
 143                     spell_menu.addAction(action)
 144                 # Only add the spelling suggests to the menu if there are
 145                 # suggestions.
 146                 if len(spell_menu.actions()) > 0:
 147                     popup_menu.addSeparator()
 148                     popup_menu.addMenu(spell_menu)
 149
 150         return popup_menu, spell_menu
 151
 152     def contextMenuEvent(self, event):
 153         popup_menu = self.context_menu()
 154         popup_menu.exec_(self.mapToGlobal(event.pos()))
 155
 156     def correct(self, word):
 157         """Replaces the selected text with word."""
 158         cursor = self.textCursor()
 159         cursor.beginEditBlock()
 160
 161         cursor.removeSelectedText()
 162         cursor.insertText(word)
 163
 164         cursor.endEditBlock()
 165
 166
 167 class Highlighter(QSyntaxHighlighter):
 168
 169     WORDS = "(?iu)[\w']+"
 170
 171     def __init__(self, doc, spellcheck):
 172         QSyntaxHighlighter.__init__(self, doc)
 173         self.spellcheck = spellcheck
 174         self.enabled = False
 175
 176     def enable(self, enabled):
 177         self.enabled = enabled
 178         self.rehighlight()
 179
 180     def highlightBlock(self, text):
 181         if not self.enabled:
 182             return
 183         text = unicode(text)
 184         fmt = QTextCharFormat()
 185         fmt.setUnderlineColor(Qt.red)
 186         fmt.setUnderlineStyle(QTextCharFormat.SpellCheckUnderline)
 187
 188         for word_object in re.finditer(self.WORDS, text):
 189             if not self.spellcheck.check(word_object.group()):
 190                 self.setFormat(word_object.start(),
 191                     word_object.end() - word_object.start(), fmt)
 192
 193
 194 class SpellAction(QAction):
 195     """QAction that returns the text in a signal.
 196     """
 197
 198     def __init__(self, *args):
 199         QAction.__init__(self, *args)
 200         self.connect(self, SIGNAL('triggered()'), self.correct)
 201
 202     def correct(self):
 203         self.emit(SIGNAL('correct'), unicode(self.text()))
 204
 205
 206 def main(args=sys.argv):
 207     app = QApplication(args)
 208
 209     widget = SpellCheckTextEdit('Type here')
 210     widget.show()
 211     widget.raise_()
 212
 213     return app.exec_()
 214
 215
 216 if __name__ == '__main__':
 217     sys.exit(main())