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

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 except TypeError as exc: 

49 raise RuntimeError("[GTR003] failed lookup of key because of unexpected exception") from exc 

50 else: 

51 raise TrieTypeError( 

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

53 tag=ErrorTag.REMOVAL_INVALID_KEY_TYPE 

54 ) 

55 

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

57 raise TrieKeyError( 

58 msg="key not found", 

59 tag=ErrorTag.REMOVAL_KEY_NOT_FOUND) 

60 

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

62 # and remove the node from the trie. 

63 node = self._trie_index[ident] 

64 del self._trie_index[ident] 

65 del self._trie_entries[ident] 

66 

67 # Remove the id from the node 

68 node.ident = None 

69 

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

71 if node.children: 

72 return 

73 

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

75 # nodes up the trie tree as needed. 

76 token: Optional[TrieKeyToken] = node.token 

77 parent = node.parent 

78 while parent is not None: 

79 del parent.children[token] 

80 # explicitly break possible cyclic references 

81 node.parent = node.token = None 

82 

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

84 if parent.ident or parent.children: 

85 return 

86 # Keep purging nodes up the tree 

87 token = parent.token 

88 node = parent 

89 parent = node.parent 

90 return