From c45aedf480c96cd10cffafb0e1518f703e2d5ae2 Mon Sep 17 00:00:00 2001
From: Dan Christensen <jdc@uwo.ca>
Date: Sun, 2 Feb 2014 23:42:10 -0500
Subject: [PATCH 1/3] Add tests for helpers.prune_split, one of which fails.

---
 attic/testsuite/helpers.py | 43 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/attic/testsuite/helpers.py b/attic/testsuite/helpers.py
index 0ff79ab76..4ad66f2a2 100644
--- a/attic/testsuite/helpers.py
+++ b/attic/testsuite/helpers.py
@@ -1,8 +1,9 @@
-from datetime import datetime
+from time import mktime, strptime
+from datetime import datetime, timezone
 import os
 import tempfile
 import unittest
-from attic.helpers import adjust_patterns, exclude_path, Location, format_timedelta, IncludePattern, ExcludePattern, make_path_safe, UpgradableLock
+from attic.helpers import adjust_patterns, exclude_path, Location, format_timedelta, IncludePattern, ExcludePattern, make_path_safe, UpgradableLock, prune_split, to_localtime
 from attic.testsuite import AtticTestCase
 
 
@@ -97,3 +98,41 @@ def test_read_only_lock_file(self):
         lock = UpgradableLock(file.name)
         self.assert_raises(UpgradableLock.LockUpgradeFailed, lock.upgrade)
         lock.release()
+
+
+class MockArchive(object):
+
+    def __init__(self, ts):
+        self.ts = ts
+
+    def __repr__(self):
+        return repr(self.ts)
+
+
+class PruneSplitTestCase(AtticTestCase):
+
+    def test(self):
+
+        def local_to_UTC(month, day):
+            'Convert noon on the month and day in 2013 to UTC.'
+            seconds = mktime(strptime('2013-%02d-%02d 12:00' % (month, day), '%Y-%m-%d %H:%M'))
+            return datetime.fromtimestamp(seconds, tz=timezone.utc)
+
+        def subset(lst, indices):
+            return {lst[i] for i in indices}
+
+        def dotest(test_archives, n, skip, indices):
+            for ta in test_archives, reversed(test_archives):
+                self.assert_equal(set(prune_split(ta, '%Y-%m', n, skip)),
+                                  subset(test_archives, indices))
+            
+        test_pairs = [(1,1), (2,1), (2,28), (3,1), (3,2), (3,31), (5,1)]
+        test_dates = [local_to_UTC(month, day) for month, day in test_pairs]
+        test_archives = [MockArchive(date) for date in test_dates]
+
+        dotest(test_archives, 3, [], [6, 5, 2])
+        dotest(test_archives, -1, [], [6, 5, 2, 0])
+        dotest(test_archives, 3, [test_archives[6]], [5, 2, 0])
+        dotest(test_archives, 3, [test_archives[5]], [6, 2, 0])
+        dotest(test_archives, 3, [test_archives[4]], [6, 5, 2])
+        dotest(test_archives, 0, [], [])

From 13863e64f8d4399ff5d0f3b9cac1d462551aeeee Mon Sep 17 00:00:00 2001
From: Dan Christensen <jdc@uwo.ca>
Date: Sun, 2 Feb 2014 23:44:04 -0500
Subject: [PATCH 2/3] Correct a theoretical bug in helpers.prune_split

---
 attic/helpers.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/attic/helpers.py b/attic/helpers.py
index 1d9b5ba69..e8a0d8354 100644
--- a/attic/helpers.py
+++ b/attic/helpers.py
@@ -99,10 +99,11 @@ def prune_split(archives, pattern, n, skip=[]):
         items.setdefault(key, [])
         items[key].append(a)
     for key, values in sorted(items.items(), reverse=True):
-        if n and values[0] not in skip:
+        if n:
             values.sort(key=attrgetter('ts'), reverse=True)
-            keep.append(values[0])
-            n -= 1
+            if values[0] not in skip:
+                keep.append(values[0])
+                n -= 1
     return keep
 
 

From d3ff6ac25dd3f875c76ba9028d818f205b7c0f92 Mon Sep 17 00:00:00 2001
From: Dan Christensen <jdc@uwo.ca>
Date: Sun, 2 Feb 2014 23:45:53 -0500
Subject: [PATCH 3/3] Clean-up helpers.prune_split

---
 attic/helpers.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/attic/helpers.py b/attic/helpers.py
index e8a0d8354..c5a7ca125 100644
--- a/attic/helpers.py
+++ b/attic/helpers.py
@@ -92,18 +92,17 @@ def write(self):
 
 
 def prune_split(archives, pattern, n, skip=[]):
-    items = {}
+    last = None
     keep = []
-    for a in archives:
-        key = to_localtime(a.ts).strftime(pattern)
-        items.setdefault(key, [])
-        items[key].append(a)
-    for key, values in sorted(items.items(), reverse=True):
-        if n:
-            values.sort(key=attrgetter('ts'), reverse=True)
-            if values[0] not in skip:
-                keep.append(values[0])
-                n -= 1
+    if n == 0:
+        return keep
+    for a in sorted(archives, key=attrgetter('ts'), reverse=True):
+        period = a.ts.strftime(pattern)
+        if period != last:
+            last = period
+            if a not in skip:
+                keep.append(a)
+                if len(keep) == n: break
     return keep