From 8e2b2be6faf21d54a6fe0e409e9f8598ab55142b Mon Sep 17 00:00:00 2001
From: Chris Barnes <barnesc@janelia.hhmi.org>
Date: Wed, 13 Nov 2019 15:26:21 +0000
Subject: [PATCH] Fix bug due to zarr constructor being function

---
 elf/io/extensions.py        | 31 +++++++++++++++++++++++++++++--
 elf/io/files.py             | 14 ++++++++++++--
 test/io_tests/test_files.py |  4 ++--
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/elf/io/extensions.py b/elf/io/extensions.py
index 35c9806..4c9a022 100644
--- a/elf/io/extensions.py
+++ b/elf/io/extensions.py
@@ -1,10 +1,13 @@
+import functools
+
 import numpy as np
 
 from .knossos_wrapper import KnossosFile, KnossosDataset
 
 
 __all__ = [
-    "FILE_CONSTRUCTORS", "GROUP_LIKE", "DATASET_LIKE", "h5py", "z5py", "pyn5", "zarr"
+    "FILE_CONSTRUCTORS", "GROUP_LIKE", "DATASET_LIKE",
+    "h5py", "z5py", "pyn5", "zarr", "zarr_open",
 ]
 
 FILE_CONSTRUCTORS = {}
@@ -61,14 +64,38 @@ def register_filetype(constructor, extensions=(), groups=(), datasets=()):
 except ImportError:
     pyn5 = None
 
+
+def identity(arg):
+    return arg
+
+
+def noop(*args, **kwargs):
+    pass
+
+
 try:
     # will not override z5py
     import zarr
+
+    # zarr stores cannot be used as context managers,
+    # which breaks compatibility with similar libraries.
+    # This wrapper patches in those methods.
+    @functools.wraps(zarr.open)
+    def zarr_open(*args, **kwargs):
+        z = zarr.open(*args, **kwargs)
+        ztype = type(z)
+        if not hasattr(ztype, "__enter__"):
+            ztype.__enter__ = identity
+        if not hasattr(ztype, "__exit__"):
+            ztype.__exit__ = noop
+        return z
+
     register_filetype(
-        zarr.open, N5_EXTS + ZARR_EXTS, zarr.hierarchy.Group, zarr.core.Array
+        zarr_open, N5_EXTS + ZARR_EXTS, zarr.hierarchy.Group, zarr.core.Array
     )
 except ImportError:
     zarr = None
+    zarr_open = None
 
 # Are there any typical knossos extensions?
 # add knossos (no extension)
diff --git a/elf/io/files.py b/elf/io/files.py
index 126d4f3..0696c52 100644
--- a/elf/io/files.py
+++ b/elf/io/files.py
@@ -1,6 +1,8 @@
 import os
-from .extensions import FILE_CONSTRUCTORS, GROUP_LIKE, DATASET_LIKE
-from .extensions import h5py, z5py
+from .extensions import (
+    FILE_CONSTRUCTORS, GROUP_LIKE, DATASET_LIKE,
+    h5py, z5py, pyn5, zarr,
+)
 from .knossos_wrapper import KnossosFile, KnossosDataset
 
 
@@ -61,6 +63,14 @@ def is_h5py(node):
     return h5py and isinstance(node, (h5py.Dataset, h5py.Group))
 
 
+def is_zarr(node):
+    return zarr and isinstance(node, (zarr.core.Array, zarr.hierarchy.Group))
+
+
+def is_pyn5(node):
+    return pyn5 and isinstance(node, (pyn5.Dataset, pyn5.Group))
+
+
 def is_knossos(node):
     """ Check if this is a KnossosWrapper object
     """
diff --git a/test/io_tests/test_files.py b/test/io_tests/test_files.py
index d74caa2..71f764e 100644
--- a/test/io_tests/test_files.py
+++ b/test/io_tests/test_files.py
@@ -4,7 +4,7 @@
 from shutil import rmtree
 
 import numpy as np
-from elf.io.extensions import h5py, z5py, pyn5, zarr, FILE_CONSTRUCTORS
+from elf.io.extensions import h5py, z5py, pyn5, zarr, zarr_open, FILE_CONSTRUCTORS
 
 
 class FileTestBase(unittest.TestCase):
@@ -81,7 +81,7 @@ class TestPyn5Files(FileTestBase, FileTestMixin):
 @unittest.skipUnless(zarr, "Need zarr")
 class TestZarrFiles(FileTestBase, FileTestMixin):
     ext = ".zr"
-    constructor = getattr(zarr, "open", None)
+    constructor = staticmethod(zarr_open)
 
 
 class TestBackendPreference(unittest.TestCase):