Add len() support to Trie

Author: bact

pythainlp/util/trie.py

@@ -16,18 +16,19 @@ def __init__(self):
             self.children = {}
 
     def __init__(self, words: Iterable[str]):
-        self.words = words
+        self.words = set(words)
         self.root = Trie.Node()
 
         for word in words:
-            self.add_word(word)
+            self.add(word)
 
-    def add_word(self, word: str) -> None:
+    def add(self, word: str) -> None:
         """
         Add a word to the trie.
 
         :param str text: a word
         """
+        self.words.add(word)
         cur = self.root
         for ch in word:
             child = cur.children.get(ch)
@@ -62,6 +63,9 @@ def __contains__(self, key: str) -> bool:
     def __iter__(self) -> Iterable[str]:
         yield from self.words
 
+    def __len__(self) -> int:
+        return len(self.words)
+
 
 def dict_trie(dict_source: Union[str, Iterable[str], Trie]) -> Trie:
     """

tests/test_util.py

@@ -250,6 +250,9 @@ def test_trie(self):
         self.assertIsNotNone(Trie(Trie(["ทดสอบ", "ทดลอง"])))
 
         trie = Trie(["ทด", "ทดสอบ", "ทดลอง"])
+        self.assertTrue("ทด" in trie)
+        trie.add("ทบ")
+        self.assertEqual(len(trie), 4)
         self.assertEqual(len(trie.prefixes("ทดสอบ")), 2)
 
         self.assertIsNotNone(dict_trie(Trie(["ลอง", "ลาก"])))