util.py: add DictOfSets (with tests), LRUCache, defaultdict (for py2.4)

buildbot/test/unit/test_util.py

@@ -54,3 +54,23 @@ def test_file_path(self):
     def test_win32file_path(self):
         self.assertUrl('c:\\repos\\my-repo', 'c:\\repos\\my-repo')
 
+class DictStuff(unittest.TestCase):
+    def test_dictofsets(self):
+        d = util.DictOfSets()
+        d.add("a", 1)
+        d.add("a", 2)
+        d.add("b", 3)
+        self.failUnless("a" in d)
+        self.failUnless("b" in d)
+        self.failIf("c" in d)
+        self.failUnlessEqual(d["a"], set([1,2]))
+        self.failUnlessEqual(d["b"], set([3]))
+        self.failUnlessRaises(KeyError, lambda: d["c"])
+        d.remove("c", 4) # ignored
+        d.remove("b", 4) # ignored
+        d.remove("b", 3)
+        self.failIf("b" in d)
+        self.failUnlessEqual(d.pop("a"), set([1,2]))
+        self.failIf("a" in d)
+        self.failUnlessEqual(d.pop("c"), set())
+        self.failUnlessEqual(d.pop("f"), set())

buildbot/util.py

@@ -2,6 +2,7 @@
 
 from twisted.internet.defer import Deferred
 from twisted.spread import pb
+from twisted.python import threadable
 import time, re, string
 
 def naturalSort(l):
@@ -131,3 +132,63 @@ def remove_userpassword(url):
     repo_url = protocol_url[1].split('@')[-1]
 
     return protocol + '://' + repo_url
+
+class LRUCache:
+    synchronized = ["get", "add"]
+
+    def __init__(self, max_size=50):
+        self._max_size = max_size
+        self._cache = {} # basic LRU cache
+        self._cached_ids = [] # = [LRU .. MRU]
+
+    def get(self, id):
+        thing = self._cache.get(id, None)
+        if thing:
+            self._cached_ids.remove(id)
+            self._cached_ids.append(id)
+        return thing
+
+    def add(self, id, thing):
+        if id in self._cache:
+            return
+        while len(self._cached_ids) >= self._max_size:
+            del self._cache[self._cached_ids.pop(0)]
+        self._cache[id] = thing
+        self._cached_ids.append(id)
+
+threadable.synchronize(LRUCache)
+
+
+# collections.defaultdict only appeared in py2.5, but buildbot supports 2.4
+class defaultdict(dict):
+    def __init__(self, default_factory=None, *args, **kwargs):
+        self._default_factory = default_factory
+        dict.__init__(self, *args, **kwargs)
+    def __getitem__(self, key):
+        if key not in self and self._default_factory:
+            self[key] = self._default_factory()
+        return dict.__getitem__(self, key)
+
+class DictOfSets:
+    # a bit like defaultdict(set), but don't keep empty sets around, so it
+    # doesn't grow forever. "key in d" can be used to rule out empty sets.
+    # Also don't create a set just to probe for members.
+    def __init__(self):
+        self.d = dict()
+    def add(self, key, value):
+        if key not in self.d:
+            self.d[key] = set()
+        self.d[key].add(value)
+    def remove(self, key, value):
+        if key in self.d:
+            self.d[key].discard(value)
+            if not self.d[key]:
+                del self.d[key]
+    def __contains__(self, key):
+        return key in self.d
+    def __getitem__(self, key):
+        return self.d[key]
+    def pop(self, key):
+        if key in self.d:
+            return self.d.pop(key)
+        return set()