Slice items can now be a contextlib.AbstractContextManager; deprecated SliceSource.dispose() and introduced SliceSource.close().

Author: forman

tests/slice/test_cm.py

@@ -2,6 +2,7 @@
 # Permissions are hereby granted under the terms of the MIT License:
 # https://opensource.org/licenses/MIT.
 
+import contextlib
 import shutil
 import unittest
 import warnings
@@ -13,6 +14,7 @@
 from zappend.fsutil.fileobj import FileObj
 from zappend.slice.cm import SliceSourceContextManager
 from zappend.slice.cm import open_slice_dataset
+from zappend.slice.source import SliceSource
 from zappend.slice.sources.memory import MemorySliceSource
 from zappend.slice.sources.persistent import PersistentSliceSource
 from zappend.slice.sources.temporary import TemporarySliceSource
@@ -23,12 +25,11 @@
 # noinspection PyUnusedLocal
 
 
-# noinspection PyShadowingBuiltins
+# noinspection PyShadowingBuiltins,PyRedeclaration
 class OpenSliceDatasetTest(unittest.TestCase):
     def setUp(self):
         clear_memory_fs()
 
-    # noinspection PyMethodMayBeStatic
     def test_slice_item_is_slice_source(self):
         dataset = make_test_dataset()
         ctx = Context(dict(target_dir="memory://target.zarr"))
@@ -127,7 +128,6 @@ def test_slice_item_is_uri_with_polling_ok(self):
         with slice_cm as slice_ds:
             self.assertIsInstance(slice_ds, xr.Dataset)
 
-    # noinspection PyMethodMayBeStatic
     def test_slice_item_is_uri_with_polling_fail(self):
         slice_dir = FileObj("memory://slice.zarr")
         ctx = Context(
@@ -140,3 +140,114 @@ def test_slice_item_is_uri_with_polling_fail(self):
         with pytest.raises(FileNotFoundError, match=slice_dir.uri):
             with slice_cm:
                 pass
+
+    def test_slice_item_is_context_manager(self):
+        @contextlib.contextmanager
+        def get_dataset(name):
+            uri = f"memory://{name}.zarr"
+            ds = make_test_dataset(uri=uri)
+            try:
+                yield ds
+            finally:
+                ds.close()
+                FileObj(uri).delete(recursive=True)
+
+        ctx = Context(
+            dict(
+                target_dir="memory://target.zarr",
+                slice_source=get_dataset,
+            )
+        )
+        slice_cm = open_slice_dataset(ctx, "bibo")
+        self.assertIsInstance(slice_cm, contextlib.AbstractContextManager)
+        with slice_cm as slice_ds:
+            self.assertIsInstance(slice_ds, xr.Dataset)
+
+    def test_slice_item_is_slice_source(self):
+        class MySliceSource(SliceSource):
+            def __init__(self, name):
+                self.uri = f"memory://{name}.zarr"
+                self.ds = None
+
+            def get_dataset(self):
+                self.ds = make_test_dataset(uri=self.uri)
+                return self.ds
+
+            def close(self):
+                if self.ds is not None:
+                    self.ds.close()
+                FileObj(uri=self.uri).delete(recursive=True)
+
+        ctx = Context(
+            dict(
+                target_dir="memory://target.zarr",
+                slice_source=MySliceSource,
+            )
+        )
+        slice_cm = open_slice_dataset(ctx, "bibo")
+        self.assertIsInstance(slice_cm, SliceSourceContextManager)
+        self.assertIsInstance(slice_cm.slice_source, SliceSource)
+        with slice_cm as slice_ds:
+            self.assertIsInstance(slice_ds, xr.Dataset)
+
+    def test_slice_item_is_deprecated_slice_source(self):
+        class MySliceSource(SliceSource):
+            def __init__(self, name):
+                self.uri = f"memory://{name}.zarr"
+                self.ds = None
+
+            def get_dataset(self):
+                self.ds = make_test_dataset(uri=self.uri)
+                return self.ds
+
+            def dispose(self):
+                if self.ds is not None:
+                    self.ds.close()
+                FileObj(uri=self.uri).delete(recursive=True)
+
+        ctx = Context(
+            dict(
+                target_dir="memory://target.zarr",
+                slice_source=MySliceSource,
+            )
+        )
+        slice_cm = open_slice_dataset(ctx, "bibo")
+        self.assertIsInstance(slice_cm, SliceSourceContextManager)
+        self.assertIsInstance(slice_cm.slice_source, SliceSource)
+        with pytest.warns(expected_warning=DeprecationWarning):
+            with slice_cm as slice_ds:
+                self.assertIsInstance(slice_ds, xr.Dataset)
+
+
+class IsContextManagerTest(unittest.TestCase):
+    """Assert that context managers are identified by isinstance()"""
+
+    def test_context_manager_class(self):
+        @contextlib.contextmanager
+        def my_slice_source(data):
+            ds = xr.Dataset(data)
+            try:
+                yield ds
+            finally:
+                ds.close()
+
+        item = my_slice_source([1, 2, 3])
+        self.assertTrue(isinstance(item, contextlib.AbstractContextManager))
+        self.assertFalse(isinstance(my_slice_source, contextlib.AbstractContextManager))
+
+    def test_context_manager_protocol(self):
+        class MySliceSource:
+            def __enter__(self):
+                return xr.Dataset()
+
+            def __exit__(self, *exc):
+                pass
+
+        item = MySliceSource()
+        self.assertTrue(isinstance(item, contextlib.AbstractContextManager))
+        self.assertFalse(isinstance(MySliceSource, contextlib.AbstractContextManager))
+
+    def test_dataset(self):
+        item = xr.Dataset()
+        self.assertTrue(isinstance(item, contextlib.AbstractContextManager))
+        self.assertFalse(isinstance(xr.Dataset, contextlib.AbstractContextManager))

tests/slice/test_source.py

@@ -2,24 +2,19 @@
 # Permissions are hereby granted under the terms of the MIT License:
 # https://opensource.org/licenses/MIT.
 
-import shutil
 import unittest
-import warnings
 
 import pytest
 import xarray as xr
 
 from zappend.context import Context
 from zappend.fsutil.fileobj import FileObj
-from zappend.slice.cm import SliceSourceContextManager
-from zappend.slice.cm import open_slice_dataset
-from zappend.slice.source import to_slice_source, SliceSource
+from zappend.slice.source import SliceSource
+from zappend.slice.source import to_slice_source
 from zappend.slice.sources.memory import MemorySliceSource
 from zappend.slice.sources.persistent import PersistentSliceSource
 from zappend.slice.sources.temporary import TemporarySliceSource
 from tests.helpers import clear_memory_fs
-from tests.helpers import make_test_dataset
-from tests.config.test_config import CustomSliceSource
 
 
 # noinspection PyUnusedLocal
@@ -86,22 +81,43 @@ def my_slice_source(arg1, arg2=None, ctx=None):
             return xr.Dataset(attrs=dict(arg1=arg1, arg2=arg2, ctx=ctx))
 
         ctx = make_ctx(slice_source=my_slice_source)
-        arg = xr.Dataset()
         slice_source = to_slice_source(ctx, ([13], {"arg2": True}), 0)
         self.assertIsInstance(slice_source, MemorySliceSource)
         ds = slice_source.get_dataset()
         self.assertEqual(13, ds.attrs.get("arg1"))
         self.assertEqual(True, ds.attrs.get("arg2"))
         self.assertIs(ctx, ds.attrs.get("ctx"))
 
+    def test_slice_item_is_slice_source_context_manager(self):
+        import contextlib
+
+        @contextlib.contextmanager
+        def my_slice_source(ctx, arg1, arg2=None):
+            _ds = xr.Dataset(attrs=dict(arg1=arg1, arg2=arg2, ctx=ctx))
+            try:
+                yield _ds
+            finally:
+                _ds.close()
+
+        ctx = make_ctx(slice_source=my_slice_source)
+        slice_source = to_slice_source(ctx, ([14], {"arg2": "OK"}), 0)
+        self.assertIsInstance(slice_source, contextlib.AbstractContextManager)
+        with slice_source as ds:
+            self.assertIsInstance(ds, xr.Dataset)
+            self.assertEqual(14, ds.attrs.get("arg1"))
+            self.assertEqual("OK", ds.attrs.get("arg2"))
+            self.assertIs(ctx, ds.attrs.get("ctx"))
+
     # noinspection PyMethodMayBeStatic
     def test_raises_if_slice_item_is_int(self):
         ctx = make_ctx(persist_mem_slices=True)
         with pytest.raises(
             TypeError,
             match=(
                 "slice_item must have type str, xarray.Dataset,"
-                " zappend.api.FileObj, zappend.api.SliceSource, but was type int"
+                " contextlib.AbstractContextManager,"
+                " zappend.api.FileObj, zappend.api.SliceSource,"
+                " but was type int"
             ),
         ):
             to_slice_source(ctx, 42, 0)
@@ -116,7 +132,9 @@ def hallo():
             TypeError,
             match=(
                 "slice_item must have type str, xarray.Dataset,"
-                " zappend.api.FileObj, zappend.api.SliceSource, but was type function"
+                " contextlib.AbstractContextManager,"
+                " zappend.api.FileObj, zappend.api.SliceSource,"
+                " but was type function"
             ),
         ):
             to_slice_source(ctx, hallo, 0)

tests/test_api.py

@@ -150,7 +150,7 @@ def __init__(self, slice_ds):
             def get_dataset(self) -> xr.Dataset:
                 return self.slice_ds.drop_vars(["tsm"])
 
-            def dispose(self):
+            def close(self):
                 pass
 
         target_dir = "memory://target.zarr"
@@ -168,6 +168,28 @@ def dispose(self):
             ds.attrs,
         )
 
+    def test_some_slices_with_slice_source_cm(self):
+        import contextlib
+
+        @contextlib.contextmanager
+        def drop_tsm(slice_ds: xr.Dataset) -> xr.Dataset:
+            yield slice_ds.drop_vars(["tsm"])
+
+        target_dir = "memory://target.zarr"
+        slices = [make_test_dataset(index=3 * i) for i in range(3)]
+        zappend(slices, target_dir=target_dir, slice_source=drop_tsm)
+        ds = xr.open_zarr(target_dir)
+        self.assertEqual({"time": 9, "y": 50, "x": 100}, ds.sizes)
+        self.assertEqual({"chl"}, set(ds.data_vars))
+        self.assertEqual({"time", "y", "x"}, set(ds.coords))
+        self.assertEqual(
+            {
+                "Conventions": "CF-1.8",
+                "title": "Test 1-3",
+            },
+            ds.attrs,
+        )
+
     def test_some_slices_with_slice_source_func(self):
         def drop_tsm(slice_ds: xr.Dataset) -> xr.Dataset:
             return slice_ds.drop_vars(["tsm"])

zappend/slice/cm.py

@@ -3,7 +3,7 @@
 # https://opensource.org/licenses/MIT.
 
 import contextlib
-from typing import Any
+from typing import Any, ContextManager
 
 import xarray as xr
 
@@ -15,6 +15,8 @@
 class SliceSourceContextManager(contextlib.AbstractContextManager):
     """A context manager that wraps a slice source.
 
+    Internal class, no API.
+
     Args:
         slice_source: The slice source.
     """
@@ -30,15 +32,15 @@ def __enter__(self) -> xr.Dataset:
         return self._slice_source.get_dataset()
 
     def __exit__(self, *exception_args):
-        self._slice_source.dispose()
+        self._slice_source.close()
         self._slice_source = None
 
 
 def open_slice_dataset(
     ctx: Context,
     slice_item: Any,
     slice_index: int = 0,
-) -> SliceSourceContextManager:
+) -> ContextManager[xr.Dataset]:
     """Open the slice source for given slice item `slice_item`.
 
     The intended and only use of the returned slice source is as context
@@ -77,4 +79,7 @@ class derived from `zappend.slice.SliceSource` or a function that returns
         A new slice source instance
     """
     slice_source = to_slice_source(ctx, slice_item, slice_index)
-    return SliceSourceContextManager(slice_source)
+    if isinstance(slice_source, contextlib.AbstractContextManager):
+        return slice_source
+    else:
+        return SliceSourceContextManager(slice_source)