From bfda3532ea58a48533ceaa417b1bd5c3f5137be3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89tienne=20Simon?= Date: Mon, 11 May 2015 21:20:11 +0000 Subject: Add method to TaxiDataset to easily extract a dictionary --- data/hdf5.py | 5 +++++ 1 file changed, 5 insertions(+) (limited to 'data/hdf5.py') diff --git a/data/hdf5.py b/data/hdf5.py index d848023..6e2f9a4 100644 --- a/data/hdf5.py +++ b/data/hdf5.py @@ -19,6 +19,11 @@ class TaxiDataset(H5PYDataset): def data_path(self): return os.path.join(data.path, self.filename) + def extract(self, request): + if not self.load_in_memory: + raise ValueError('extract called on a dataset not loaded in memory') + return dict(zip(self.sources, self.get_data(None, request))) + class TaxiStream(DataStream): def __init__(self, which_set, filename='data.hdf5', iteration_scheme=None, **kwargs): dataset = TaxiDataset(which_set, filename, **kwargs) -- cgit v1.2.3