From 94fbea6266e263bd2650399a5d2aed3540552f07 Mon Sep 17 00:00:00 2001
From: Valentin Lorentz <progval@progval.net>
Date: Thu, 31 Jul 2014 14:24:47 +0200
Subject: [PATCH] Add utils.structures.TruncatableSet.

---
 src/utils/structures.py | 35 +++++++++++++++++++++++++++++++++++
 test/test_utils.py      | 22 +++++++++++++++++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/src/utils/structures.py b/src/utils/structures.py
index 31e8f561f..1be386ffd 100644
--- a/src/utils/structures.py
+++ b/src/utils/structures.py
@@ -454,5 +454,40 @@ class CacheDict(collections.MutableMapping):
     def __len__(self):
         return len(self.d)
 
+class TruncatableSet(collections.MutableSet):
+    """A set that keeps track of the order of inserted elements so
+    the oldest can be removed."""
+    def __init__(self, iterable):
+        self._ordered_items = list(iterable)
+        self._items = set(self._ordered_items)
+    def __contains__(self, item):
+        return item in self._items
+    def __iter__(self):
+        return iter(self._items)
+    def __len__(self):
+        return len(self._items)
+    def add(self, item):
+        if item not in self._items:
+            self._items.add(item)
+            self._ordered_items.append(item)
+    def discard(self, item):
+        self._items.discard(item)
+        self._ordered_items.remove(item)
+    def truncate(self, size):
+        assert size >= 0
+        removed_size = len(self)-size
+        # I make two different cases depending on removed_size<size
+        # in order to make if faster if one is significantly bigger than the
+        # other.
+        if removed_size <= 0:
+            return
+        elif removed_size < size:
+            # If there are more kept items than removed items
+            for old_item in self._ordered_items[0:-size]:
+                self.discard(old_item)
+            self._ordered_items = self._ordered_items[-size:]
+        else:
+            self._ordered_items = self._ordered_items[-size:]
+            self._items = set(self._ordered_items)
 
 # vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79:
diff --git a/test/test_utils.py b/test/test_utils.py
index bffb0d0b0..7032a69f6 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -1132,7 +1132,27 @@ class TestCacheDict(SupyTestCase):
             self.failUnless(len(d) <= max)
             self.failUnless(i in d)
             self.failUnless(d[i] == i)
-            
+
+class TestTruncatableSet(SupyTestCase):
+    def testBasics(self):
+        s = TruncatableSet(['foo', 'bar', 'baz', 'qux'])
+        self.assertEqual(s, {'foo', 'bar', 'baz', 'qux'})
+        self.failUnless('foo' in s)
+        self.failUnless('bar' in s)
+        self.failIf('quux' in s)
+        s.discard('baz')
+        self.failUnless('foo' in s)
+        self.failIf('baz' in s)
+        s.add('quux')
+        self.failUnless('quux' in s)
+
+    def testTruncate(self):
+        s = TruncatableSet(['foo', 'bar'])
+        s.add('baz')
+        s.add('qux')
+        s.truncate(3)
+        self.assertEqual(s, {'bar', 'baz', 'qux'})
+
 
 # vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79: