From 922359ce8e266e057c40a3cd47089f66ac0f6b74 Mon Sep 17 00:00:00 2001
From: Andy Salnikov <salnikov@slac.stanford.edu>
Date: Tue, 14 Nov 2023 10:19:16 -0800
Subject: [PATCH] Add `Butler._caching_context` method.

Registry shiim now redirects to this method instead of Registry
method. RemoteButler raises NotImplementedError but may do something
non-trivial later when we know how caching is going to work with
client/server.
---
 python/lsst/daf/butler/_butler.py                      | 5 +++++
 python/lsst/daf/butler/_registry_shim.py               | 6 ++----
 python/lsst/daf/butler/direct_butler.py                | 4 ++++
 python/lsst/daf/butler/registry/_registry.py           | 3 +--
 python/lsst/daf/butler/remote_butler/_remote_butler.py | 6 ++++++
 5 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/python/lsst/daf/butler/_butler.py b/python/lsst/daf/butler/_butler.py
index f83a4ad347..937ff8e6c9 100644
--- a/python/lsst/daf/butler/_butler.py
+++ b/python/lsst/daf/butler/_butler.py
@@ -482,6 +482,11 @@ def get_known_repos(cls) -> set[str]:
         """
         return ButlerRepoIndex.get_known_repos()
 
+    @abstractmethod
+    def _caching_context(self) -> AbstractContextManager[None]:
+        """Context manager that enables caching."""
+        raise NotImplementedError()
+
     @abstractmethod
     def transaction(self) -> AbstractContextManager[None]:
         """Context manager supporting `Butler` transactions.
diff --git a/python/lsst/daf/butler/_registry_shim.py b/python/lsst/daf/butler/_registry_shim.py
index 2cce959eba..4d2653abe0 100644
--- a/python/lsst/daf/butler/_registry_shim.py
+++ b/python/lsst/daf/butler/_registry_shim.py
@@ -102,11 +102,9 @@ def refresh(self) -> None:
         # Docstring inherited from a base class.
         self._registry.refresh()
 
-    @contextlib.contextmanager
-    def caching_context(self) -> Iterator[None]:
+    def caching_context(self) -> contextlib.AbstractContextManager[None]:
         # Docstring inherited from a base class.
-        with self._registry.caching_context():
-            yield
+        return self._butler._caching_context()
 
     @contextlib.contextmanager
     def transaction(self, *, savepoint: bool = False) -> Iterator[None]:
diff --git a/python/lsst/daf/butler/direct_butler.py b/python/lsst/daf/butler/direct_butler.py
index 6b70ecb1e1..80bcf5d1ae 100644
--- a/python/lsst/daf/butler/direct_butler.py
+++ b/python/lsst/daf/butler/direct_butler.py
@@ -299,6 +299,10 @@ def isWriteable(self) -> bool:
         # Docstring inherited.
         return self._registry.isWriteable()
 
+    def _caching_context(self) -> contextlib.AbstractContextManager[None]:
+        """Context manager that enables caching."""
+        return self._registry.caching_context()
+
     @contextlib.contextmanager
     def transaction(self) -> Iterator[None]:
         """Context manager supporting `Butler` transactions.
diff --git a/python/lsst/daf/butler/registry/_registry.py b/python/lsst/daf/butler/registry/_registry.py
index 444a67eb54..21e651314c 100644
--- a/python/lsst/daf/butler/registry/_registry.py
+++ b/python/lsst/daf/butler/registry/_registry.py
@@ -118,9 +118,8 @@ def refresh(self) -> None:
         """
         raise NotImplementedError()
 
-    @contextlib.contextmanager
     @abstractmethod
-    def caching_context(self) -> Iterator[None]:
+    def caching_context(self) -> contextlib.AbstractContextManager[None]:
         """Context manager that enables caching."""
         raise NotImplementedError()
 
diff --git a/python/lsst/daf/butler/remote_butler/_remote_butler.py b/python/lsst/daf/butler/remote_butler/_remote_butler.py
index 841735a28b..1930461ec0 100644
--- a/python/lsst/daf/butler/remote_butler/_remote_butler.py
+++ b/python/lsst/daf/butler/remote_butler/_remote_butler.py
@@ -158,6 +158,12 @@ def _simplify_dataId(
         # Assume we can treat it as a dict.
         return SerializedDataCoordinate(dataId=data_id)
 
+    def _caching_context(self) -> AbstractContextManager[None]:
+        # Docstring inherited.
+        # Not implemented for now, will have to think whether this needs to
+        # do something on client side and/or remote side.
+        raise NotImplementedError()
+
     def transaction(self) -> AbstractContextManager[None]:
         """Will always raise NotImplementedError.
         Transactions are not supported by RemoteButler.