Merge pull request #1530 from quantopian/add-specific-assets

ENH: Add `SpecificAssets` filter.

tests/pipeline/test_filter.py

@@ -26,11 +26,13 @@ from numpy import (
 from numpy.random import randn, seed as random_seed
 
 from zipline.errors import BadPercentileBounds
-from zipline.pipeline import Filter, Factor
+from zipline.pipeline import Filter, Factor, Pipeline
 from zipline.pipeline.classifiers import Classifier
 from zipline.pipeline.factors import CustomFactor
-from zipline.pipeline.filters import All, Any, AtLeastN
-from zipline.testing import parameter_space, permute_rows
+from zipline.pipeline.filters import All, Any, AtLeastN, SpecificAssets
+from zipline.testing import parameter_space, permute_rows, ZiplineTestCase
+from zipline.testing.fixtures import WithSeededRandomPipelineEngine
+from zipline.testing.predicates import assert_equal
 from zipline.utils.numpy_utils import float64_dtype, int64_dtype
 from .base import BasePipelineTestCase, with_default_shape
 
@@ -820,3 +822,40 @@ class FilterTestCase(BasePipelineTestCase):
             },
             mask=self.build_mask(permute(rot90(self.eye_mask(shape=shape)))),
         )
+
+
+class SpecificAssetsTestCase(WithSeededRandomPipelineEngine,
+                             ZiplineTestCase):
+
+    ASSET_FINDER_EQUITY_SIDS = tuple(range(10))
+
+    def test_specific_assets(self):
+        assets = self.asset_finder.retrieve_all(self.ASSET_FINDER_EQUITY_SIDS)
+
+        class SidFactor(CustomFactor):
+            """A factor that just returns each asset's sid."""
+            inputs = ()
+            window_length = 1
+
+            def compute(self, today, sids, out):
+                out[:] = sids
+
+        pipe = Pipeline(
+            columns={
+                'sid': SidFactor(),
+                'evens': SpecificAssets(assets[::2]),
+                'odds': SpecificAssets(assets[1::2]),
+                'first_five': SpecificAssets(assets[:5]),
+                'last_three': SpecificAssets(assets[-3:]),
+            },
+        )
+
+        start, end = self.trading_days[[-10, -1]]
+        results = self.run_pipeline(pipe, start, end).unstack()
+
+        sids = results.sid.astype(int64_dtype)
+
+        assert_equal(results.evens, ~(sids % 2).astype(bool))
+        assert_equal(results.odds, (sids % 2).astype(bool))
+        assert_equal(results.first_five, sids < 5)
+        assert_equal(results.last_three, sids >= 7)