def load_corpus(corpus, case_sensitive=False):
"""Builds a Trie from a list of strings."""
root = TrieNode()
for word in corpus:
if case_sensitive:
root.insert(word, word)
else:
root.insert(word.lower(), word)
return root
class TestTrieNode(unittest.TestCase):
def setUp(self):
self.root = TrieNode()
self.root.insert('ab', 'ab')
self.root.insert('bc', 'bc')
self.root.insert('cd', 'cd')
self.root.insert('bcde', 'bcde')
self.root.insert('bcdf', 'teste')
def test_insert_root(self):
self.assertEqual(len(self.root.children), 3)
def test_children(self):
node = self.root.children['a']
self.assertEqual(len(node.children), 1)
self.assertEqual(node.value, None)
node = self.root.children['b']
self.assertEqual(len(node.children), 1)
self.assertEqual(node.value, None)
node = node.children['c']
self.assertEqual(len(node.children), 1)
self.assertEqual(node.value, 'bc')
node = node.children['d']
self.assertEqual(len(node.children), 2)
self.assertEqual(node.value, None)
node_e = node.children['e']
self.assertEqual(len(node_e.children), 0)
self.assertEqual(node_e.value, 'bcde')
node_f = node.children['f']
self.assertEqual(len(node_f.children), 0)
self.assertEqual(node_f.value, 'teste')
def test_find_completions_node(self):
node = find_completions_node(self.root, 'a')
self.assertEqual(node, self.root.children['a'])
node = find_completions_node(self.root, 'b')
self.assertEqual(node, self.root.children['b'])
node = find_completions_node(self.root, 'c')
self.assertEqual(node, self.root.children['c'])
node = find_completions_node(self.root, 'c')
self.assertEqual(node, self.root.children['c'])
node = find_completions_node(self.root, 'ab')
self.assertEqual(node, self.root.children['a'].children['b'])
node = find_completions_node(self.root, 'bc')
self.assertEqual(node, self.root.children['b'].children['c'])
node = find_completions_node(self.root, 'cd')
self.assertEqual(node, self.root.children['c'].children['d'])
node = find_completions_node(self.root, 'bcd')
self.assertEqual(node,
self.root.children['b'].children['c'].children['d'])
node = find_completions_node(self.root, 'bcde')
self.assertEqual(
node,
self.root.children['b'].children['c'].children['d'].children['e'])
node = find_completions_node(self.root, 'bcdf')
self.assertEqual(
node,
self.root.children['b'].children['c'].children['d'].children['f'])