Improve the code that temporarily changes the formatter location

Now has a context manager rather than directly changing the
internal property of the formatter.

python/lsst/daf/butler/core/formatter.py

@@ -25,9 +25,22 @@
 
 from abc import ABCMeta, abstractmethod
 from collections.abc import Mapping
+import contextlib
 import logging
 import copy
-from typing import ClassVar, Set, AbstractSet, Union, Optional, Dict, Any, Tuple, Type, TYPE_CHECKING
+from typing import (
+    AbstractSet,
+    Any,
+    ClassVar,
+    Dict,
+    Iterator,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    TYPE_CHECKING,
+    Union,
+)
 
 from .configSupport import processLookupConfigs, LookupKey
 from .mappingFactory import MappingFactory
@@ -267,6 +280,40 @@ def toBytes(self, inMemoryDataset: Any) -> bytes:
         """
         raise NotImplementedError("Type does not support writing to bytes.")
 
+    @contextlib.contextmanager
+    def _updateLocation(self, location: Location) -> Iterator[Location]:
+        """Temporarily replace the location associated with this formatter.
+
+        Parameters
+        ----------
+        location : `Location`
+            New location to use for this formatter. If `None` the
+            formatter will not change but it will still return
+            the old location. This allows it to be used in a code
+            path where the location may not need to be updated
+            but the with block is still convenient.
+
+        Yields
+        ------
+        old : `Location`
+            The old location that will be restored.
+
+        Notes
+        -----
+        This is an internal method that should be used with care.
+        It may change in the future. Should be used as a context
+        manager to restore the location when the temporary is no
+        longer required.
+        """
+        old = self._fileDescriptor.location
+        try:
+            if location is not None:
+                self._fileDescriptor.location = location
+            yield old
+        finally:
+            if location is not None:
+                self._fileDescriptor.location = old
+
     def makeUpdatedLocation(self, location: Location) -> Location:
         """Return a new `Location` instance updated with this formatter's
         extension.

python/lsst/daf/butler/datastores/fileLikeDatastore.py

@@ -874,9 +874,9 @@ def _removeFileExists(uri: ButlerURI) -> None:
                     # Need to configure the formatter to write to a different
                     # location and that needs us to overwrite internals
                     tmpLocation = Location(*os.path.split(tmpFile.name))
-                    formatter._fileDescriptor.location = tmpLocation
                     log.debug("Writing dataset to temporary location at %s", tmpLocation.uri)
-                    formatter.write(inMemoryDataset)
+                    with formatter._updateLocation(tmpLocation):
+                        formatter.write(inMemoryDataset)
                     uri.transfer_from(tmpLocation.uri, transfer="copy", overwrite=True)
                 log.debug("Successfully wrote dataset to %s via a temporary file.", uri)
 
@@ -941,15 +941,17 @@ def _read_artifact_into_memory(self, getInfo: DatastoreFileGetInformation,
                 # because formatter.read does not allow an override.
                 # This could be improved.
                 msg = ""
+                newLocation = None
                 if uri != local_uri:
-                    formatter._fileDescriptor.location = Location(*local_uri.split())
+                    newLocation = Location(*local_uri.split())
                     msg = "(via download to local file)"
 
                 log.debug("Reading %s from location %s %s with formatter %s",
                           f"component {getInfo.component}" if isComponent else "",
                           uri, msg, formatter.name())
                 try:
-                    result = formatter.read(component=getInfo.component if isComponent else None)
+                    with formatter._updateLocation(newLocation):
+                        result = formatter.read(component=getInfo.component if isComponent else None)
                 except Exception as e:
                     raise ValueError(f"Failure from formatter '{formatter.name()}' for dataset {ref.id}"
                                      f" ({ref.datasetType.name} from {uri}): {e}") from e