From a2cab5f5e170a297322bc8e0d0f85bef5786509a Mon Sep 17 00:00:00 2001
From: Jeremy Magland <jmagland@flatironinstitute.org>
Date: Thu, 21 Mar 2024 07:28:08 -0400
Subject: [PATCH] properly close resources

---
 lindi/LindiH5ZarrStore/LindiH5ZarrStore.py | 22 +++++++++++-----------
 lindi/LindiH5pyFile/LindiH5pyFile.py       |  4 +++-
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py b/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py
index 5f87aa1..b6a23f2 100644
--- a/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py
+++ b/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py
@@ -52,15 +52,17 @@ def __init__(
         _file: Union[IO, Any],
         _opts: LindiH5ZarrStoreOpts,
         _url: Union[str, None] = None,
+        _entities_to_close: List[Any]
     ):
         """
         Do not call the constructor directly. Instead, use the from_file class
         method.
         """
-        self._file = _file
-        self._h5f = h5py.File(_file, "r")
+        self._file: Union[IO, Any, None] = _file
+        self._h5f: Union[h5py.File, None] = h5py.File(_file, "r")
         self._url = _url
         self._opts = _opts
+        self._entities_to_close = _entities_to_close + [self._h5f]
 
         # Some datasets do not correspond to traditional chunked datasets. For
         # those datasets, we need to store the inline data so that we can return
@@ -97,21 +99,19 @@ def from_file(
         if hdf5_file_name_or_url.startswith(
             "http://"
         ) or hdf5_file_name_or_url.startswith("https://"):
+            # note that the remfile.File object does not need to be closed
             remf = remfile.File(hdf5_file_name_or_url, verbose=False)
-            return LindiH5ZarrStore(_file=remf, _url=hdf5_file_name_or_url, _opts=opts)
+            return LindiH5ZarrStore(_file=remf, _url=hdf5_file_name_or_url, _opts=opts, _entities_to_close=[])
         else:
             f = open(hdf5_file_name_or_url, "rb")
-            return LindiH5ZarrStore(_file=f, _url=url, _opts=opts)
+            return LindiH5ZarrStore(_file=f, _url=url, _opts=opts, _entities_to_close=[f])
 
     def close(self):
         """Close the store."""
-        if hasattr(self, "_h5f") and self._h5f:
-            self._h5f.close()
-            self._h5f = None
-        if hasattr(self, "_file") and self._file:
-            if not isinstance(self._file, remfile.File):
-                self._file.close()
-            self._file = None
+        for e in self._entities_to_close:
+            e.close()
+        self._h5f = None
+        self._file = None
 
     def __getitem__(self, key):
         """Get an item from the store (required by base class)."""
diff --git a/lindi/LindiH5pyFile/LindiH5pyFile.py b/lindi/LindiH5pyFile/LindiH5pyFile.py
index 93b39b7..41a3074 100644
--- a/lindi/LindiH5pyFile/LindiH5pyFile.py
+++ b/lindi/LindiH5pyFile/LindiH5pyFile.py
@@ -54,6 +54,8 @@ def from_zarr_store(zarr_store: Union[ZarrStore, FSMap]):
         """
         Create a LindiH5pyFile from a zarr store.
         """
+        # note that even though the function is called "open", the zarr_group
+        # does not need to be closed
         zarr_group = zarr.open(store=zarr_store, mode="r")
         assert isinstance(zarr_group, zarr.Group)
         return LindiH5pyFile.from_zarr_group(zarr_group)
@@ -122,7 +124,7 @@ def swmr_mode(self, value):  # type: ignore
         raise Exception("Getting swmr_mode is not allowed")
 
     def close(self):
-        # return self._file_object.close()
+        # Nothing was opened, so nothing needs to be closed
         pass
 
     def flush(self):