Merge branch '8791' of https://github.com/vikram/pandas into vikram-8791

Author: jreback

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

doc/source/whatsnew/v0.15.2.txt

@@ -77,6 +77,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`` for automatic closing (:issue:`8791`).
 
 .. _whatsnew_0152.performance:
 

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

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: