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'