diff --git a/README.md b/README.md index d2bf0b5..0a6e6db 100644 --- a/README.md +++ b/README.md @@ -20,8 +20,7 @@ from retrie.trie import Trie trie = Trie() -for term in ["abc", "foo", "abs"]: - trie.add(term) +trie.add("abc", "foo", "abs") assert trie.pattern() == "(?:ab[cs]|foo)" # equivalent to but faster than "(?:abc|abs|foo)" trie.add("absolute") diff --git a/src/retrie/retrie.py b/src/retrie/retrie.py index 38a2d42..1047534 100644 --- a/src/retrie/retrie.py +++ b/src/retrie/retrie.py @@ -123,8 +123,7 @@ def __init__( Retrie.__init__(self, word_boundary=word_boundary, re_flags=re_flags) - for term in keys: - self.trie.add(term) + self.trie.add(*keys) @cached_property def compiled(self): # type: (...) -> Pattern[Text] @@ -291,6 +290,6 @@ def replace( Args: text (str): String to search & replace. - count (int): Amount of occurences to replace. If 0 or emitted, replace all. + count (int): Amount of occurences to replace. If 0 or omitted, replace all. """ return self.compiled.sub(self._replace, text, count=count) diff --git a/src/retrie/trie.py b/src/retrie/trie.py index e8843e6..7b5a8a4 100644 --- a/src/retrie/trie.py +++ b/src/retrie/trie.py @@ -17,13 +17,14 @@ class Trie: def __init__(self): self.data = {} # type: data_type - def add(self, word): - """Add a word to the current Trie.""" - ref = self.data - for char in word: - ref[char] = ref.get(char, {}) - ref = ref[char] - ref[""] = 1 + def add(self, *word): + """Add one or more words to the current Trie.""" + for word in word: + ref = self.data + for char in word: + ref[char] = ref.get(char, {}) + ref = ref[char] + ref[""] = 1 def dump(self): # type: (...) -> data_type """Dump the current trie as dictionary.""" diff --git a/tests/test_trie.py b/tests/test_trie.py index 2bc78bc..4483001 100644 --- a/tests/test_trie.py +++ b/tests/test_trie.py @@ -5,8 +5,7 @@ def test_Trie(): trie = Trie() assert trie.pattern() == "" - for term in ["abc", "foo", "abs"]: - trie.add(term) + trie.add("abc", "foo", "abs") assert trie.pattern() == "(?:ab[cs]|foo)" trie.add("absolute")