diff --git a/spellchecker/spellchecker.py b/spellchecker/spellchecker.py index 7bccbba..4a30d55 100644 --- a/spellchecker/spellchecker.py +++ b/spellchecker/spellchecker.py @@ -20,9 +20,12 @@ class SpellChecker(object): for no dictionary. Supported languages are `en`, `es`, `de`, and \ `fr`. Defaults to `en` local_dictionary (str): The path to a locally stored word \ - frequency dictionary ''' + frequency dictionary + distance (int): The edit distance to use. Defaults to 2''' - def __init__(self, language='en', local_dictionary=None): + + def __init__(self, language='en', local_dictionary=None, distance=2): + self._distance = distance self._word_frequency = WordFrequency() if local_dictionary: self._word_frequency.load_dictionary(local_dictionary) @@ -95,8 +98,9 @@ def candidates(self, word): word (str): The word for which to calculate candidate spellings Returns: set: The set of words that are possible candidates ''' + return (self.known([word]) or self.known(self.edit_distance_1(word)) or - self.known(self.edit_distance_2(word)) or {word}) + (self._distance == 2 and self.known(self.edit_distance_2(word))) or {word}) def known(self, words): ''' The subset of `words` that appear in the dictionary of words diff --git a/tests/resources/small_dictionary.json b/tests/resources/small_dictionary.json index 8b62a44..df0a4ca 100644 --- a/tests/resources/small_dictionary.json +++ b/tests/resources/small_dictionary.json @@ -1,5 +1,6 @@ { "a": 1, "b": 2, - "apple": 45 + "apple": 45, + "bike": 60 } diff --git a/tests/spellchecker_test.py b/tests/spellchecker_test.py index adde440..48c450c 100644 --- a/tests/spellchecker_test.py +++ b/tests/spellchecker_test.py @@ -6,7 +6,6 @@ from spellchecker import SpellChecker - class TestSpellChecker(unittest.TestCase): ''' test the spell checker class ''' @@ -117,6 +116,13 @@ def test_load_external_dictionary(self): self.assertEqual(spell['a'], 1) self.assertTrue('apple' in spell) + def test_edit_distance_one(self): + ''' test a case where edit distance must be one ''' + here = os.path.dirname(__file__) + filepath = '{}/resources/small_dictionary.json'.format(here) + spell = SpellChecker(language=None, local_dictionary=filepath, distance=1) + self.assertEqual(spell.candidates('hike'), {'bike'}) + def test_edit_distance_two(self): ''' test a case where edit distance must be two ''' here = os.path.dirname(__file__)