Coverage for src/gentrie/trie/removal.py: 100%

1# -*- coding: utf-8 -*- 

2"""Entry removal operations for the trie.""" 

3from typing import Optional 

4 

5from ..exceptions import ErrorTag, TrieKeyError, TrieTypeError 

6from ..protocols import GeneralizedKey, TrieKeyToken 

7from ..types import TrieId 

8from ..validation import is_generalizedkey 

9 

10from .trie_mixins import TrieMixinsInterface 

11 

12# Disabled because pyright does not understand mixins 

13# use of private attributes from the mixing class as declared 

14# in the TrieMixinsInterface protocol. 

15# pyright: reportPrivateUsage=false 

16 

17 

18class TrieRemovalMixin: 

19 """Mixin providing entry removal operations. 

20 

21 Note: This mixin accesses private attributes of the mixing class. 

22 This is intentional and necessary for the mixin pattern 

23 """ 

24 

25 def remove(self: TrieMixinsInterface, key: TrieId | GeneralizedKey) -> None: 

26 """Remove the specified key from the trie. 

27 

28 Removes the key from the trie. If the key is not found, it raises a KeyError. 

29 The key can be specified either as a :class:`TrieId` or as a :class:`GeneralizedKey`. 

30 

31 Args: 

32 key (TrieId | GeneralizedKey): identifier for key to remove. 

33 

34 Raises: 

35 TrieTypeError: if the key arg is not a :class:`TrieId` or a valid :class:`GeneralizedKey`. 

36 TrieKeyError: if the key arg does not match the id or trie key of any entries in the trie. 

37 """ 

38 ident: Optional[TrieId] = None 

39 if isinstance(key, TrieId): 

40 ident = key 

41 # If runtime validation is disabled, we just ASSUME the key is a GeneralizedKey 

42 # if we got to here. 

43 elif (not self.runtime_validation) or is_generalizedkey(key): 

44 try: 

45 ident = self[key].ident 

46 except KeyError: 

47 ident = None 

48 else: 

49 raise TrieTypeError( 

50 msg="key arg must be of type TrieId or a valid GeneralizedKey", 

51 tag=ErrorTag.REMOVAL_INVALID_KEY_TYPE 

52 ) 

53 

54 if ident is None or ident not in self._trie_index: 

55 raise TrieKeyError( 

56 msg="key not found", 

57 tag=ErrorTag.REMOVAL_KEY_NOT_FOUND) 

58 

59 # Get the node and delete its id from the trie index and entries 

60 # and remove the node from the trie. 

61 node = self._trie_index[ident] 

62 del self._trie_index[ident] 

63 del self._trie_entries[ident] 

64 

65 # Remove the id from the node 

66 node.ident = None 

67 

68 # If the node still has other trie ids or children, we're done: return 

69 if node.children: 

70 return 

71 

72 # No trie ids or children are left for this node, so prune 

73 # nodes up the trie tree as needed. 

74 token: Optional[TrieKeyToken] = node.token 

75 parent = node.parent 

76 while parent is not None: 

77 del parent.children[token] 

78 # explicitly break possible cyclic references 

79 node.parent = node.token = None 

80 

81 # If the parent node has a trie id or children, we're done: return 

82 if parent.ident or parent.children: 

83 return 

84 # Keep purging nodes up the tree 

85 token = parent.token 

86 node = parent 

87 parent = node.parent 

88 return