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

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

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

3 from typing import Optional

4

5 from ..exceptions import ErrorTag, TrieKeyError, TrieTypeError

6 from ..protocols import GeneralizedKey, TrieKeyToken

7 from ..types import TrieId

8 from ..validation import is_generalizedkey

9

10 from .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

18 class 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