BUG: Fixes should_clean for keep_last=0

volconst · Aug 17, 2016 · fcf1067 · fcf1067
1 parent 440806a
commit fcf1067
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 4 deletions.
diff --git a/tests/data/bundles/test_core.py b/tests/data/bundles/test_core.py
@@ -7,7 +7,7 @@
 
 from zipline.assets.synthetic import make_simple_equity_info
 from zipline.data.bundles import UnknownBundle, from_bundle_ingest_dirname
-from zipline.data.bundles.core import _make_bundle_core
+from zipline.data.bundles.core import _make_bundle_core, BadClean
 from zipline.lib.adjustment import Float64Multiply
 from zipline.pipeline.loaders.synthetic import (
     make_bar_data,
@@ -369,6 +369,26 @@ def test_clean_keep_last(self):
             msg='keep_last=2 did not remove the correct number of ingestions',
         )
 
+        with assert_raises(BadClean):
+            self.clean('bundle', keep_last=-1, environ=self.environ)
+
+        assert_equal(
+            self._list_bundle(),
+            {fourth, fifth},
+            msg='keep_last=-1 removed some ingestions',
+        )
+
+        assert_equal(
+            self.clean('bundle', keep_last=0, environ=self.environ),
+            {fourth, fifth},
+        )
+
+        assert_equal(
+            self._list_bundle(),
+            set(),
+            msg='keep_last=0 did not remove the correct number of ingestions',
+        )
+
     @staticmethod
     def _ts_of_run(run):
         return from_bundle_ingest_dirname(run.rsplit(os.path.sep, 1)[-1])

diff --git a/zipline/data/bundles/core.py b/zipline/data/bundles/core.py
@@ -7,7 +7,7 @@
 from contextlib2 import ExitStack
 import click
 import pandas as pd
-from toolz import curry, complement
+from toolz import curry, complement, take
 
 from ..us_equity_pricing import (
     BcolzDailyBarReader,
@@ -563,11 +563,13 @@ def should_clean(name):
                     (after is not None and dt > after)
                 )
 
-        else:
-            last_n_dts = set(all_runs[-keep_last:])
+        elif keep_last >= 0:
+            last_n_dts = set(take(keep_last, reversed(all_runs)))
 
             def should_clean(name):
                 return name not in last_n_dts
+        else:
+            raise BadClean(before, after, keep_last)
 
         cleaned = set()
         for run in all_runs: