From f2dc979fbc276cc1dcc0c33dae3d7c4e29a7cdb9 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Tue, 17 Dec 2013 11:13:12 -0500 Subject: [PATCH] BUG: Make all iteration related methods fo BarData match __iter__ `for s in data` and methods like `for s in data.keys` were not producing the same list of active sids Make the other iteration methods match __iter__ by using the contains method to check whether or not the sid is active. For use of data outside of the algoscript context, which needs access to all data fields use data._data --- zipline/protocol.py | 29 ++++++++++++++++++--------- zipline/transforms/batch_transform.py | 4 ++-- 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/zipline/protocol.py b/zipline/protocol.py index cc3f0e59..4834fb64 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -164,23 +164,32 @@ class BarData(object): def __iter__(self): for sid, data in self._data.iteritems(): - if len(data): - yield sid - - def keys(self): - return self._data.keys() + # Allow contains override to filter out sids. + if sid in self: + if len(data): + yield sid def iterkeys(self): - return self._data.iterkeys() + # Allow contains override to filter out sids. + return [sid for sid in self._data.iterkeys() if sid in self] + + def keys(self): + # Allow contains override to filter out sids. + return list(self.iterkeys()) def itervalues(self): - return self._data.itervalues() + return (value for sid, value in self.iteritems()) + + def values(self): + return list(self.itervalues()) def iteritems(self): - return self._data.iteritems() + return ((sid, value) for sid, value + in self._data.iteritems() + if sid in self) def items(self): - return self._data.items() + return list(self.iteritems()) def __len__(self): - return len(self._data.keys()) + return len(self.keys()) diff --git a/zipline/transforms/batch_transform.py b/zipline/transforms/batch_transform.py index 5f0698ea..41bfa203 100644 --- a/zipline/transforms/batch_transform.py +++ b/zipline/transforms/batch_transform.py @@ -230,7 +230,7 @@ class BatchTransform(object): Point of entry. Process an event frame. """ # extract dates - dts = [event.datetime for event in data.itervalues()] + dts = [event.datetime for event in data._data.itervalues()] # we have to provide the event with a dt. This is only for # checking if the event is outside the window or not so a # couple of seconds shouldn't matter. We don't add it to @@ -238,7 +238,7 @@ class BatchTransform(object): # sid keys. event = Event() event.dt = max(dts) - event.data = {k: v.__dict__ for k, v in data.iteritems() + event.data = {k: v.__dict__ for k, v in data._data.iteritems() # Need to check if data has a 'length' to filter # out sids without trade data available. # TODO: expose more of 'no trade available'