diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 521fa9f5..985f363c 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -15,7 +15,11 @@ from logbook import Logger, Processor import zipline.finance.trading as trading -from zipline.protocol import BarData, DATASOURCE_TYPE +from zipline.protocol import ( + BarData, + SIDData, + DATASOURCE_TYPE +) from zipline.gens.utils import hash_args log = Logger('Trade Simulation') @@ -198,5 +202,8 @@ class AlgorithmSimulator(object): Update the universe with new event information. """ # Update our knowledge of this event's sid - sid_data = self.current_data[event.sid] + if event.sid in self.current_data: + sid_data = self.current_data[event.sid] + else: + sid_data = self.current_data[event.sid] = SIDData() sid_data.__dict__.update(event.__dict__) diff --git a/zipline/protocol.py b/zipline/protocol.py index 0376b49a..b32314eb 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -12,7 +12,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from collections import defaultdict import datetime from . utils.protocol_utils import Enum @@ -135,7 +134,7 @@ class BarData(object): """ def __init__(self): - self._data = defaultdict(SIDData) + self._data = {} self._contains_override = None def __contains__(self, name): @@ -164,7 +163,9 @@ class BarData(object): del self._data[name] def __iter__(self): - return self._data.iterkeys() + for sid, data in self._data.iteritems(): + if len(data): + yield sid def keys(self): return self._data.keys()