diff --git a/doc/source/io.rst b/doc/source/io.rst index ba19173546cb2..f2d5924edac77 100644 --- a/doc/source/io.rst +++ b/doc/source/io.rst @@ -2348,7 +2348,7 @@ Closing a Store, Context Manager # Working with, and automatically closing the store with the context # manager - with get_store('store.h5') as store: + with HDFStore('store.h5') as store: store.keys() .. ipython:: python diff --git a/doc/source/whatsnew/v0.15.2.txt b/doc/source/whatsnew/v0.15.2.txt index 260e4e8a17c7d..57534fa2e1ef0 100644 --- a/doc/source/whatsnew/v0.15.2.txt +++ b/doc/source/whatsnew/v0.15.2.txt @@ -75,6 +75,7 @@ Enhancements - Added ``Timedelta.to_timedelta64`` method to the public API (:issue:`8884`). - Added ``gbq.generate_bq_schema`` function to the gbq module (:issue:`8325`). - ``Series`` now works with map objects the same way as generators (:issue:`8909`). +- Added context manager to ``HDFStore`` (:issue:`8791`). Now with HDFStore('path') as store: should also close to store. get_store will be deprecated in the future. .. _whatsnew_0152.performance: diff --git a/pandas/io/pytables.py b/pandas/io/pytables.py index 56c444095ca51..05510f655f7be 100644 --- a/pandas/io/pytables.py +++ b/pandas/io/pytables.py @@ -251,33 +251,6 @@ def _tables(): return _table_mod -@contextmanager -def get_store(path, **kwargs): - """ - Creates an HDFStore instance. This function can be used in a with statement - - Parameters - ---------- - same as HDFStore - - Examples - -------- - >>> from pandas import DataFrame - >>> from numpy.random import randn - >>> bar = DataFrame(randn(10, 4)) - >>> with get_store('test.h5') as store: - ... store['foo'] = bar # write to HDF5 - ... bar = store['foo'] # retrieve - """ - store = None - try: - store = HDFStore(path, **kwargs) - yield store - finally: - if store is not None: - store.close() - - # interface to/from ### def to_hdf(path_or_buf, key, value, mode=None, complevel=None, complib=None, @@ -289,7 +262,7 @@ def to_hdf(path_or_buf, key, value, mode=None, complevel=None, complib=None, f = lambda store: store.put(key, value, **kwargs) if isinstance(path_or_buf, string_types): - with get_store(path_or_buf, mode=mode, complevel=complevel, + with HDFStore(path_or_buf, mode=mode, complevel=complevel, complib=complib) as store: f(store) else: @@ -493,6 +466,12 @@ def __unicode__(self): return output + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + def keys(self): """ Return a (potentially unordered) list of the keys corresponding to the @@ -1288,6 +1267,12 @@ def _read_group(self, group, **kwargs): return s.read(**kwargs) +def get_store(path, **kwargs): + """ Backwards compatible alias for ``HDFStore`` + """ + return HDFStore(path, **kwargs) + + class TableIterator(object): """ define the iteration interface on a table diff --git a/pandas/io/tests/test_pytables.py b/pandas/io/tests/test_pytables.py index e2a99076e7607..e95d46f66f17f 100644 --- a/pandas/io/tests/test_pytables.py +++ b/pandas/io/tests/test_pytables.py @@ -172,6 +172,25 @@ def test_factory_fun(self): finally: safe_remove(self.path) + def test_context(self): + try: + with HDFStore(self.path) as tbl: + raise ValueError('blah') + except ValueError: + pass + finally: + safe_remove(self.path) + + try: + with HDFStore(self.path) as tbl: + tbl['a'] = tm.makeDataFrame() + + with HDFStore(self.path) as tbl: + self.assertEqual(len(tbl), 1) + self.assertEqual(type(tbl['a']), DataFrame) + finally: + safe_remove(self.path) + def test_conv_read_write(self): try: @@ -334,10 +353,10 @@ def test_api_default_format(self): pandas.set_option('io.hdf.default_format','table') df.to_hdf(path,'df3') - with get_store(path) as store: + with HDFStore(path) as store: self.assertTrue(store.get_storer('df3').is_table) df.to_hdf(path,'df4',append=True) - with get_store(path) as store: + with HDFStore(path) as store: self.assertTrue(store.get_storer('df4').is_table) pandas.set_option('io.hdf.default_format',None) @@ -463,11 +482,11 @@ def check(mode): # context if mode in ['r','r+']: def f(): - with get_store(path,mode=mode) as store: + with HDFStore(path,mode=mode) as store: pass self.assertRaises(IOError, f) else: - with get_store(path,mode=mode) as store: + with HDFStore(path,mode=mode) as store: self.assertEqual(store._handle.mode, mode) with ensure_clean_path(self.path) as path: