1- from typing import Dict
1+ # Copyright 2016-2021 The Matrix.org Foundation C.I.C.
2+ #
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+ #
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+ #
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
214
315SENTINEL = object ()
416
517
18+ class TreeCacheNode (dict ):
19+ """The type of nodes in our tree.
20+
21+ Has its own type so we can distinguish it from real dicts that are stored at the
22+ leaves.
23+ """
24+
25+ pass
26+
27+
628class TreeCache :
729 """
830 Tree-based backing store for LruCache. Allows subtrees of data to be deleted
931 efficiently.
1032 Keys must be tuples.
33+
34+ The data structure is a chain of TreeCacheNodes:
35+ root = {key_1: {key_2: _value}}
1136 """
1237
1338 def __init__ (self ):
1439 self .size = 0
15- self .root = {} # type: Dict
40+ self .root = TreeCacheNode ()
1641
1742 def __setitem__ (self , key , value ):
1843 return self .set (key , value )
@@ -21,10 +46,23 @@ def __contains__(self, key):
2146 return self .get (key , SENTINEL ) is not SENTINEL
2247
2348 def set (self , key , value ):
49+ if isinstance (value , TreeCacheNode ):
50+ # this would mean we couldn't tell where our tree ended and the value
51+ # started.
52+ raise ValueError ("Cannot store TreeCacheNodes in a TreeCache" )
53+
2454 node = self .root
2555 for k in key [:- 1 ]:
26- node = node .setdefault (k , {})
27- node [key [- 1 ]] = _Entry (value )
56+ next_node = node .get (k , SENTINEL )
57+ if next_node is SENTINEL :
58+ next_node = node [k ] = TreeCacheNode ()
59+ elif not isinstance (next_node , TreeCacheNode ):
60+ # this suggests that the caller is not being consistent with its key
61+ # length.
62+ raise ValueError ("value conflicts with an existing subtree" )
63+ node = next_node
64+
65+ node [key [- 1 ]] = value
2866 self .size += 1
2967
3068 def get (self , key , default = None ):
@@ -33,25 +71,41 @@ def get(self, key, default=None):
3371 node = node .get (k , None )
3472 if node is None :
3573 return default
36- return node .get (key [- 1 ], _Entry ( default )). value
74+ return node .get (key [- 1 ], default )
3775
3876 def clear (self ):
3977 self .size = 0
40- self .root = {}
78+ self .root = TreeCacheNode ()
4179
4280 def pop (self , key , default = None ):
81+ """Remove the given key, or subkey, from the cache
82+
83+ Args:
84+ key: key or subkey to remove.
85+ default: value to return if key is not found
86+
87+ Returns:
88+ If the key is not found, 'default'. If the key is complete, the removed
89+ value. If the key is partial, the TreeCacheNode corresponding to the part
90+ of the tree that was removed.
91+ """
92+ # a list of the nodes we have touched on the way down the tree
4393 nodes = []
4494
4595 node = self .root
4696 for k in key [:- 1 ]:
4797 node = node .get (k , None )
48- nodes .append (node ) # don't add the root node
4998 if node is None :
5099 return default
100+ if not isinstance (node , TreeCacheNode ):
101+ # we've gone off the end of the tree
102+ raise ValueError ("pop() key too long" )
103+ nodes .append (node ) # don't add the root node
51104 popped = node .pop (key [- 1 ], SENTINEL )
52105 if popped is SENTINEL :
53106 return default
54107
108+ # working back up the tree, clear out any nodes that are now empty
55109 node_and_keys = list (zip (nodes , key ))
56110 node_and_keys .reverse ()
57111 node_and_keys .append ((self .root , None ))
@@ -61,14 +115,15 @@ def pop(self, key, default=None):
61115
62116 if n :
63117 break
118+ # found an empty node: remove it from its parent, and loop.
64119 node_and_keys [i + 1 ][0 ].pop (k )
65120
66- popped , cnt = _strip_and_count_entires ( popped )
121+ cnt = sum ( 1 for _ in iterate_tree_cache_entry ( popped ) )
67122 self .size -= cnt
68123 return popped
69124
70125 def values (self ):
71- return list ( iterate_tree_cache_entry (self .root ) )
126+ return iterate_tree_cache_entry (self .root )
72127
73128 def __len__ (self ):
74129 return self .size
@@ -78,36 +133,9 @@ def iterate_tree_cache_entry(d):
78133 """Helper function to iterate over the leaves of a tree, i.e. a dict of that
79134 can contain dicts.
80135 """
81- if isinstance (d , dict ):
136+ if isinstance (d , TreeCacheNode ):
82137 for value_d in d .values ():
83138 for value in iterate_tree_cache_entry (value_d ):
84139 yield value
85140 else :
86- if isinstance (d , _Entry ):
87- yield d .value
88- else :
89- yield d
90-
91-
92- class _Entry :
93- __slots__ = ["value" ]
94-
95- def __init__ (self , value ):
96- self .value = value
97-
98-
99- def _strip_and_count_entires (d ):
100- """Takes an _Entry or dict with leaves of _Entry's, and either returns the
101- value or a dictionary with _Entry's replaced by their values.
102-
103- Also returns the count of _Entry's
104- """
105- if isinstance (d , dict ):
106- cnt = 0
107- for key , value in d .items ():
108- v , n = _strip_and_count_entires (value )
109- d [key ] = v
110- cnt += n
111- return d , cnt
112- else :
113- return d .value , 1
141+ yield d
0 commit comments