1from __future__ import absolute_import, division, unicode_literals
2
3from datrie import Trie as DATrie
4from six import text_type
5
6from ._base import Trie as ABCTrie
7
8
9class Trie(ABCTrie):
10    def __init__(self, data):
11        chars = set()
12        for key in data.keys():
13            if not isinstance(key, text_type):
14                raise TypeError("All keys must be strings")
15            for char in key:
16                chars.add(char)
17
18        self._data = DATrie("".join(chars))
19        for key, value in data.items():
20            self._data[key] = value
21
22    def __contains__(self, key):
23        return key in self._data
24
25    def __len__(self):
26        return len(self._data)
27
28    def __iter__(self):
29        raise NotImplementedError()
30
31    def __getitem__(self, key):
32        return self._data[key]
33
34    def keys(self, prefix=None):
35        return self._data.keys(prefix)
36
37    def has_keys_with_prefix(self, prefix):
38        return self._data.has_keys_with_prefix(prefix)
39
40    def longest_prefix(self, prefix):
41        return self._data.longest_prefix(prefix)
42
43    def longest_prefix_item(self, prefix):
44        return self._data.longest_prefix_item(prefix)
45