def initialize(context): pipeline = Pipeline() context.vwaps = [] for length in vwaps: name = vwap_key(length) factor = VWAP(window_length=length) context.vwaps.append(factor) pipeline.add(factor, name=name) filter_ = (USEquityPricing.close.latest > 300) pipeline.add(filter_, 'filter') if set_screen: pipeline.set_screen(filter_) attach_pipeline(pipeline, 'test')
def test_set_screen(self): f, g = SomeFilter(), SomeOtherFilter() p = Pipeline() self.assertEqual(p.screen, None) p.set_screen(f) self.assertEqual(p.screen, f) with self.assertRaises(ValueError): p.set_screen(f) p.set_screen(g, overwrite=True) self.assertEqual(p.screen, g) with self.assertRaises(TypeError) as e: p.set_screen(f, g) message = e.exception.args[0] self.assertIn( "expected a value of type bool or int for argument 'overwrite'", message, )