[python] Conditionally apply ContextManager cleanup logic

apis/python/src/tiledbsoma/_soma_group.py

@@ -34,6 +34,7 @@ class _CachedElement:
     uri: str
     tiledb_type: SOMABaseTileDBType | None = None
     soma: SOMAObject | None = None
+    managed: bool = False
     """The reified object, if it has been opened."""
 
     @classmethod
@@ -91,6 +92,7 @@ def __getitem__(self, key: str) -> CollectionElementType:
             if entry.soma is None:
                 from . import _factory  # Delayed binding to resolve circular import.
 
+                entry.managed = True
                 entry.soma = _factory._open_soma_object(
                     uri=entry.uri,
                     mode=self.mode,
@@ -99,8 +101,9 @@ def __getitem__(self, key: str) -> CollectionElementType:
                     clib_type=None if entry.tiledb_type is None else entry.tiledb_type.name,
                 )
 
-                # Since we just opened this object, we own it and should close it.
-                self._close_stack.enter_context(entry.soma)
+                if self._is_running_in_context:
+                    # Since we just opened this object, we own it and should close it.
+                    self._close_stack.enter_context(entry.soma)
 
         return cast("CollectionElementType", entry.soma)
 
@@ -346,6 +349,14 @@ def set(
         self._set_element(key, uri=uri_to_add, relative=use_relative_uri, soma_object=value)
         return self
 
+    def close(self, recursive: bool = True) -> None:
+        if recursive:
+            for [_, value] in self._contents.items():
+                if value.soma is not None and value.managed:
+                    value.soma.close(recursive)
+
+        super().close(recursive)
+
 
 @attrs.define(frozen=True, kw_only=True)
 class _ChildURI:

apis/python/src/tiledbsoma/_soma_object.py

@@ -31,7 +31,7 @@ class SOMAObject:
     _handle_type: ClassVar[_tdb_handles.RawHandle]
     """Class variable of the clib class handle used to open this object type."""
 
-    __slots__ = ("_close_stack", "_context", "_handle", "_metadata", "_timestamp_ms", "_uri")
+    __slots__ = ("_close_stack", "_context", "_handle", "_is_running_in_context", "_metadata", "_timestamp_ms", "_uri")
 
     soma_type: ClassVar[LiteralString]
     """A string describing the SOMA type of this object. This is constant.
@@ -154,6 +154,7 @@ def __init__(
         if not isinstance(handle, self._handle_type):
             raise TypeError("Internal error: Unexpected handle type {type(handle)}. Expected {self._handle_type}.")
         self._close_stack = ExitStack()
+        self._is_running_in_context = False
         """An exit stack to manage closing handles owned by this object.
 
         This is used to manage both our direct handle (in the case of simple
@@ -177,7 +178,6 @@ def __init__(
         self._uri = uri
         self._timestamp_ms = tiledb_timestamp_to_ms(self._handle.timestamp)
         self._metadata = _tdb_handles.MetadataWrapper.from_handle(self._handle)
-        self._close_stack.enter_context(self._handle)
         self._parse_special_metadata()
 
     def _parse_special_metadata(self) -> None:
@@ -234,13 +234,16 @@ def metadata(self) -> MutableMapping[str, Any]:
     __hash__ = object.__hash__
 
     def __enter__(self) -> Self:
+        self._is_running_in_context = True
+        self._close_stack.enter_context(self._handle)
         return self
 
     def __exit__(self, *_: Any) -> None:  # noqa: ANN401
+        self._is_running_in_context = False
         self.close()
 
     def __del__(self) -> None:
-        self.close()
+        self.close(False)
         super_del = getattr(super(), "__del__", lambda: None)
         super_del()
 
@@ -265,7 +268,7 @@ def uri(self) -> str:
         """
         return self._uri
 
-    def close(self) -> None:
+    def close(self, recursive: bool = True) -> None:
         """Release any resources held while the object is open.
         Closing an already-closed object is a no-op.
 
@@ -275,8 +278,11 @@ def close(self) -> None:
         Lifecycle:
             Maturing.
         """
+        del recursive  # unused by default
+
         if not self.closed:
             self._metadata._write()
+            self._handle.close()
         self._close_stack.close()
 
     @property

apis/python/tests/test_collection.py

@@ -669,3 +669,37 @@ def get_obsm(exp, key) -> soma.SparseNDArray:
     with soma.open(uri) as exp:
         results = list(tp.map(get_obsm, repeat(exp), ("X_umap", "X_umap", "X_umap", "X_umap")))
     assert all(r is results[0] for r in results)
+
+
+def test_close_behavior() -> None:
+    path = ROOT_DATA_DIR / "soma-experiment-versions-2025-04-04/1.16.1/pbmc3k_processed"
+    uri = str(path)
+    if not pathlib.Path(uri).is_dir():
+        raise RuntimeError(
+            f"Missing '{uri}' directory. Try running `make data` from the TileDB-SOMA project root directory.",
+        )
+
+    # When not running under a context closing the parent collection should close the child members it owns
+    exp = soma.open(uri)
+    obs = exp.obs
+    exp.close()
+    assert exp.closed
+    assert obs.closed
+
+    # When not running under a context closing the parent collection should not close the child members it owns if recusrsive is set to False
+    exp = soma.open(uri)
+    obs = exp.obs
+    exp.close(False)
+    assert exp.closed
+    assert not obs.closed
+
+    # Temporary objects will by default close with recursive set to False
+    obs = soma.open(uri).obs
+    assert not obs.closed
+
+    # When running under a context closing the parent collection should close the child members
+    with soma.open(uri) as exp:
+        obs = exp.obs
+
+    assert exp.closed
+    assert obs.closed