From 824788f55c30f8762f2c0e0d1a87f50e1ec250e5 Mon Sep 17 00:00:00 2001
From: Chris Sheppard <chris.sheppard@points.com>
Date: Fri, 27 Apr 2018 13:42:54 -0400
Subject: [PATCH] Fix small eviction bug; add the ability to track evictions

---
 expiringdict/__init__.py   | 17 ++++++++++++++---
 tests/expiringdict_test.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/expiringdict/__init__.py b/expiringdict/__init__.py
index 852f8af..648c7ca 100755
--- a/expiringdict/__init__.py
+++ b/expiringdict/__init__.py
@@ -27,13 +27,14 @@
 
 
 class ExpiringDict(OrderedDict):
-    def __init__(self, max_len, max_age_seconds):
+    def __init__(self, max_len, max_age_seconds, eviction_counter=None):
         assert max_age_seconds >= 0
         assert max_len >= 1
 
         OrderedDict.__init__(self)
         self.max_len = max_len
         self.max_age = max_age_seconds
+        self.evictions = eviction_counter
         self.lock = RLock()
 
         if sys.version_info >= (3, 5):
@@ -74,12 +75,15 @@ def __getitem__(self, key, with_age=False):
     def __setitem__(self, key, value):
         """ Set d[key] to value. """
         with self.lock:
-            if len(self) == self.max_len:
+            OrderedDict.__setitem__(self, key, (value, time.time()))
+            original_overage = len(self) - self.max_len
+            while len(self) > self.max_len:
                 try:
                     self.popitem(last=False)
                 except KeyError:
                     pass
-            OrderedDict.__setitem__(self, key, (value, time.time()))
+            if self.evictions is not None:
+                self.evictions += original_overage
 
     def pop(self, key, default=None):
         """ Get item from the dict and remove it.
@@ -137,6 +141,13 @@ def values(self):
                 pass
         return r
 
+    def reset_evictions(self, new_value=0):
+        """ Reset the eviction count.  Returns the current value """
+        with self.lock:
+            evictions = self.evictions
+            self.evictions = new_value
+        return evictions
+
     def fromkeys(self):
         " Create a new dictionary with keys from seq and values set to value. "
         raise NotImplementedError()
diff --git a/tests/expiringdict_test.py b/tests/expiringdict_test.py
index f2d0855..0fda5a3 100755
--- a/tests/expiringdict_test.py
+++ b/tests/expiringdict_test.py
@@ -64,6 +64,14 @@ def test_repr():
     sleep(0.01)
     eq_(str(d), "ExpiringDict([])")
 
+def test_eviction_on_wrong_count():
+    d = ExpiringDict(max_len=3, max_age_seconds=20)
+    d['a'] = 'x'
+    d['b'] = 'y'
+    d['c'] = 'z'
+    d['b'] = 'y'
+    eq_([k for k in d], ['a', 'b', 'c'])
+    eq_([k for k in d.values()], ['x', 'y', 'z'])
 
 def test_iter():
     d = ExpiringDict(max_len=10, max_age_seconds=0.01)
@@ -113,6 +121,35 @@ def test_setdefault():
     eq_('y', d.setdefault('a', 'y'))
 
 
+def test_eviction_counting():
+    d = ExpiringDict(max_len=1, max_age_seconds=100, eviction_counter=0)
+
+    d['a'] = 'A'
+    d['b'] = 'B'
+    d['c'] = 'C'
+
+    eq_(len(d), 1)
+    eq_(d.evictions, 2)
+
+
+def test_eviction_reset():
+    d = ExpiringDict(max_len=1, max_age_seconds=100, eviction_counter=0)
+
+    d['a'] = 'A'
+    d['b'] = 'B'
+    d['c'] = 'C'
+
+    current_eviction_count = d.reset_evictions()
+    eq_(len(d), 1)
+    eq_(current_eviction_count, 2)
+    eq_(d.evictions, 0)
+    check_eviction_count = d.reset_evictions(1) # anything that supports +=
+    d['d'] = 'D'
+    eq_(len(d), 1)
+    eq_(check_eviction_count, 0) # no operations before reset
+    eq_(d.evictions, 2)  # since we started at 1
+
+
 def test_not_implemented():
     d = ExpiringDict(max_len=10, max_age_seconds=10)
     assert_raises(NotImplementedError, d.fromkeys)