Coverage for tld/trie.py: 100%

37 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2024-05-26 22:29 +0000

1__author__ = "Artur Barseghyan" 

2__copyright__ = "2013-2023 Artur Barseghyan" 

3__license__ = "MPL-1.1 OR GPL-2.0-only OR LGPL-2.1-or-later" 

4__all__ = ( 

5 "Trie", 

6 "TrieNode", 

7) 

8 

9 

10class TrieNode(object): 

11 """Class representing a single Trie node.""" 

12 

13 __slots__ = ("children", "exception", "leaf", "private") 

14 

15 def __init__(self): 

16 self.children = None 

17 self.exception = None 

18 self.leaf = False 

19 self.private = False 

20 

21 

22class Trie(object): 

23 """An adhoc Trie data structure to store tlds in reverse notation order.""" 

24 

25 def __init__(self): 

26 self.root = TrieNode() 

27 self.__nodes = 0 

28 

29 def __len__(self): 

30 return self.__nodes 

31 

32 def add(self, tld: str, private: bool = False) -> None: 

33 node = self.root 

34 

35 # Iterating over the tld parts in reverse order 

36 # for part in reversed(tld.split('.')): 

37 tld_split = tld.split(".") 

38 tld_split.reverse() 

39 for part in tld_split: 

40 

41 if part.startswith("!"): 

42 node.exception = part[1:] 

43 break 

44 

45 # To save up some RAM, we initialize the children dict only 

46 # when strictly necessary 

47 if node.children is None: 

48 node.children = {} 

49 child = TrieNode() 

50 else: 

51 child = node.children.get(part) 

52 if child is None: 

53 child = TrieNode() 

54 

55 node.children[part] = child 

56 

57 node = child 

58 

59 node.leaf = True 

60 

61 if private: 

62 node.private = True 

63 

64 self.__nodes += 1