Skip to content

Commit

Permalink
🚸 Allow passing multiple words to Trie.add (#8)
Browse files Browse the repository at this point in the history
* 🚸 Allow passing multiple words to Trie.add

* 📝 Update docs
  • Loading branch information
ddelange authored Jul 5, 2022
1 parent 9c76a75 commit ec2bc71
Show file tree
Hide file tree
Showing 4 changed files with 12 additions and 14 deletions.
3 changes: 1 addition & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
5 changes: 2 additions & 3 deletions src/retrie/retrie.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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)
15 changes: 8 additions & 7 deletions src/retrie/trie.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
3 changes: 1 addition & 2 deletions tests/test_trie.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down

0 comments on commit ec2bc71

Please sign in to comment.