diff --git a/attic/helpers.py b/attic/helpers.py index 1d9b5ba69..c5a7ca125 100644 --- a/attic/helpers.py +++ b/attic/helpers.py @@ -92,17 +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 and values[0] not in skip: - values.sort(key=attrgetter('ts'), reverse=True) - 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 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, [], [])