diff --git a/synapse/util/caches/treecache.py b/synapse/util/caches/treecache.py
index e78305f787..c1b8ec0c73 100644
--- a/synapse/util/caches/treecache.py
+++ b/synapse/util/caches/treecache.py
@@ -64,6 +64,15 @@ class TreeCache:
self.size += 1
def get(self, key, default=None):
+ """When `key` is a full key, fetches the value for the given key (if
+ any).
+
+ If `key` is only a partial key (i.e. a truncated tuple) then returns a
+ `TreeCacheNode`, which can be passed to the `iterate_tree_cache_*`
+ functions to iterate over all entries in the cache with keys that start
+ with the given partial key.
+ """
+
node = self.root
for k in key[:-1]:
node = node.get(k, None)
@@ -139,3 +148,32 @@ def iterate_tree_cache_entry(d):
yield from iterate_tree_cache_entry(value_d)
else:
yield d
+
+
+def iterate_tree_cache_items(key, value):
+ """Helper function to iterate over the leaves of a tree, i.e. a dict of that
+ can contain dicts.
+
+ The provided key is a tuple that will get prepended to the returned keys.
+
+ Example:
+
+ cache = TreeCache()
+ cache[(1, 1)] = "a"
+ cache[(1, 2)] = "b"
+ cache[(2, 1)] = "c"
+
+ tree_node = cache.get((1,))
+
+ items = iterate_tree_cache_items((1,), tree_node)
+ assert list(items) == [((1, 1), "a"), ((1, 2), "b")]
+
+ Returns:
+ A generator yielding key/value pairs.
+ """
+ if isinstance(value, TreeCacheNode):
+ for sub_key, sub_value in value.items():
+ yield from iterate_tree_cache_items((*key, sub_key), sub_value)
+ else:
+ # we've reached a leaf of the tree.
+ yield key, value
|
