Group as context manager (#510)

For compatibility with other group-like things (e.g. `h5py.File`),
`Groups` now have enter/exit methods for use as a context manager.
If the underlying store has a `close` method, that will be called on
exit; otherwise this does nothing.

Author: clbarnes

docs/api/hierarchy.rst

@@ -11,6 +11,8 @@ Groups (``zarr.hierarchy``)
     .. automethod:: __iter__
     .. automethod:: __contains__
     .. automethod:: __getitem__
+    .. automethod:: __enter__
+    .. automethod:: __exit__
     .. automethod:: group_keys
     .. automethod:: groups
     .. automethod:: array_keys

docs/release.rst

@@ -21,7 +21,7 @@ Upcoming Release
 * Upgrade dependencies in the test matrices and resolve a
   compatibility issue with testing against the Azure Storage
   Emulator. By :user:`alimanfoo`; :issue:`468`, :issue:`467`.
-  
+
 * Do not rename Blosc parameters in n5 backend and add `blocksize` parameter,
   compatible with n5-blosc. By :user:`axtimwalde`, :issue:`485`.
 
@@ -49,6 +49,9 @@ Upcoming Release
 * Refactor out ``_tofile``/``_fromfile`` from ``DirectoryStore``.
   By :user:`John Kirkham <jakirkham>`; :issue:`503`.
 
+* Add ``__enter__``/``__exit__`` methods to ``Group`` for ``h5py.File`` compatibility.
+  By :user:`Chris Barnes <clbarnes>`; :issue:`509`.
+
 * Add documentation build to CI.
   By :user:`James Bourbeau <jrbourbeau>`; :issue:`516`.
 
@@ -97,7 +100,7 @@ Enhancements
 ~~~~~~~~~~~~
 
 * New storage backend, backed by Azure Blob Storage, class :class:`zarr.storage.ABSStore`.
-  All data is stored as block blobs. By :user:`Shikhar Goenka <shikarsg>`, 
+  All data is stored as block blobs. By :user:`Shikhar Goenka <shikarsg>`,
   :user:`Tim Crone <tjcrone>` and :user:`Zain Patel <mzjp2>`; :issue:`345`.
 
 * Add "consolidated" metadata as an experimental feature: use

docs/tutorial.rst

@@ -346,6 +346,9 @@ sub-directories, e.g.::
     >>> z
     <zarr.core.Array '/foo/bar/baz' (10000, 10000) int32>
 
+Groups can be used as context managers (in a ``with`` statement).
+If the underlying store has a ``close`` method, it will be called on exit.
+
 For more information on groups see the :mod:`zarr.hierarchy` and
 :mod:`zarr.convenience` API docs.
 

zarr/hierarchy.py

@@ -26,6 +26,8 @@ class Group(MutableMapping):
     ----------
     store : MutableMapping
         Group store, already initialized.
+        If the Group is used in a context manager, and the store has a ``close`` method,
+        it will be called on exit.
     path : string, optional
         Group path.
     read_only : bool, optional
@@ -57,6 +59,8 @@ class Group(MutableMapping):
     __iter__
     __contains__
     __getitem__
+    __enter__
+    __exit__
     group_keys
     groups
     array_keys
@@ -225,6 +229,17 @@ def __repr__(self):
         r += '>'
         return r
 
+    def __enter__(self):
+        """Return the Group for use as a context manager."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """If the underlying Store has a ``close`` method, call it."""
+        try:
+            self.store.close()
+        except AttributeError:
+            pass
+
     def info_items(self):
 
         def typestr(o):

zarr/tests/test_hierarchy.py

@@ -3,6 +3,7 @@
 import os
 import pickle
 import shutil
+import sys
 import tempfile
 import textwrap
 import unittest
@@ -860,6 +861,12 @@ def test_pickle(self):
         assert isinstance(g2['foo'], Group)
         assert isinstance(g2['foo/bar'], Array)
 
+    def test_context_manager(self):
+
+        with self.create_group() as g:
+            d = g.create_dataset('foo/bar', shape=100, chunks=10)
+            d[:] = np.arange(100)
+
 
 class TestGroupWithMemoryStore(TestGroup):
 
@@ -912,6 +919,19 @@ def create_store():
         store = ZipStore(path)
         return store, None
 
+    def test_context_manager(self):
+
+        with self.create_group() as g:
+            store = g.store
+            d = g.create_dataset('foo/bar', shape=100, chunks=10)
+            d[:] = np.arange(100)
+
+        # Check that exiting the context manager closes the store,
+        # and therefore the underlying ZipFile.
+        error = ValueError if sys.version_info >= (3, 6) else RuntimeError
+        with pytest.raises(error):
+            store.zf.extractall()
+
 
 class TestGroupWithDBMStore(TestGroup):