Merge pull request #438 from PyThaiNLP/add-dict-input-norvigspell

Add dict and other iterables support for custom_dict input

Author: bact

pythainlp/spell/pn.py

@@ -8,7 +8,17 @@
 """
 from collections import Counter
 from string import digits
-from typing import Callable, Iterable, ItemsView, List, Optional, Set, Tuple
+from typing import (
+    Callable,
+    Dict,
+    ItemsView,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
 
 from pythainlp import thai_digits, thai_letters
 from pythainlp.corpus import tnc
@@ -70,10 +80,53 @@ def _edits2(word: str) -> Set[str]:
     return set(e2 for e1 in _edits1(word) for e2 in _edits1(e1))
 
 
+def _convert_custom_dict(
+    custom_dict: Union[
+        Dict[str, int], Iterable[str], Iterable[Tuple[str, int]]
+    ],
+    min_freq: int,
+    min_len: int,
+    max_len: int,
+    dict_filter: Optional[Callable[[str], bool]],
+) -> List[Tuple[str, int]]:
+    """
+    Converts a custom dictionary to a list of (str, int) tuples
+    """
+    if isinstance(custom_dict, dict):
+        custom_dict = [(word, freq) for word, freq in custom_dict.items()]
+
+    i = iter(custom_dict)
+    first_member = next(i)
+    if isinstance(first_member, str):
+        # create tuples of a word with frequency equal to 1,
+        # and filter word list
+        custom_dict = [
+            (word, 1)
+            for word in custom_dict
+            if _keep((word, 1), 1, min_len, max_len, dict_filter)
+        ]
+    elif isinstance(first_member, tuple):
+        # filter word list
+        custom_dict = [
+            word_freq
+            for word_freq in custom_dict
+            if _keep(word_freq, min_freq, min_len, max_len, dict_filter)
+        ]
+    else:
+        raise TypeError(
+            "custom_dict must be either Dict[str, int], "
+            "Iterable[Tuple[str, int]], or Iterable[str]"
+        )
+
+    return custom_dict
+
+
 class NorvigSpellChecker:
     def __init__(
         self,
-        custom_dict: List[Tuple[str, int]] = None,
+        custom_dict: Union[
+            Dict[str, int], Iterable[str], Iterable[Tuple[str, int]]
+        ] = None,
         min_freq: int = 2,
         min_len: int = 2,
         max_len: int = 40,
@@ -91,9 +144,17 @@ def __init__(
         Then, it selects the candidate with
         the highest word occurrence probability.
 
-        :param str custom_dict: A list of tuple (word, frequency) to create
-                                a spelling dictionary. Default is from
-                                Thai National Corpus (around 40,000 words).
+        :param str custom_dict: A custom spelling dictionary. This can be:
+                                (1) a dictionary (`dict`), with words (`str`)
+                                    as keys and frequencies (`int`) as values;
+                                (2) an iterable (list, tuple, or set) of word
+                                    (`str`) and frequency (`int`) tuples:
+                                    `(str, int)`; or
+                                (3) an iterable of just words (`str`), without
+                                    frequencies -- in this case `1` will be
+                                    assigned to every words.
+                                Default is from Thai National Corpus (around
+                                40,000 words).
         :param int min_freq: Minimum frequency of a word to keep (default = 2)
         :param int min_len: Minimum length (in characters) of a word to keep
                             (default = 2)
@@ -110,12 +171,9 @@ def __init__(
         if not dict_filter:
             dict_filter = _no_filter
 
-        # filter word list
-        custom_dict = [
-            word_freq
-            for word_freq in custom_dict
-            if _keep(word_freq, min_freq, min_len, max_len, dict_filter)
-        ]
+        custom_dict = _convert_custom_dict(
+            custom_dict, min_freq, min_len, max_len, dict_filter
+        )
 
         self.__WORDS = Counter(dict(custom_dict))
         self.__WORDS += Counter()  # remove zero and negative counts

tests/test_spell.py

@@ -35,7 +35,7 @@ def test_norvig_spell_checker(self):
         self.assertTrue(len(checker.dictionary()) > 0)
         self.assertGreaterEqual(checker.prob("มี"), 0)
 
-        custom_dict = [
+        user_dict = [
             ("การงาน", 31),  # longer than max_len
             ("กาม", 1),  # fewer than min_freq
             ("กาล0", 64),  # has digit
@@ -45,6 +45,35 @@ def test_norvig_spell_checker(self):
             ("การ", 42),  # OK
         ]
         checker = NorvigSpellChecker(
-            custom_dict=custom_dict, min_freq=2, max_len=5
+            custom_dict=user_dict, min_freq=2, max_len=5
         )
         self.assertEqual(len(checker.dictionary()), 1)
+
+        user_dict = [
+            "เอกราช",
+            "ปลอดภัย",
+            "เศรษฐกิจ",
+            "เสมอภาค",
+            "เสรีภาพ",
+            "การศึกษา",
+        ]
+        checker = NorvigSpellChecker(custom_dict=user_dict)
+        self.assertEqual(len(checker.dictionary()), len(user_dict))
+
+        user_dict = {
+            "พหลโยธิน": 1,
+            "ขีตตะสังคะ": 2,
+            "พนมยงค์": 3,
+            "ภมรมนตรี": 4,
+            "มิตรภักดี": 5,
+            "ลพานุกรม": 6,
+            "สิงหเสนี": 7,
+        }
+        checker = NorvigSpellChecker(custom_dict=user_dict)
+        # "พหลโยธิน" will be removed,
+        # as it has frequency less than default min_freq (2)
+        self.assertEqual(len(checker.dictionary()), len(user_dict) - 1)
+
+        user_dict = [24, 6, 2475]
+        with self.assertRaises(TypeError):
+            checker = NorvigSpellChecker(custom_dict=user_dict)