Allow as_local to specify the temp directory to use

Author: timj

python/lsst/resources/_resourcePath.py

@@ -1011,7 +1011,7 @@ def abspath(self) -> ResourcePath:
         """
         return self
 
-    def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
+    def _as_local(self, multithreaded: bool = True, tmpdir: ResourcePath | None = None) -> tuple[str, bool]:
         """Return the location of the (possibly remote) resource as local file.
 
         This is a helper function for `as_local` context manager.
@@ -1024,6 +1024,9 @@ def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
             effect if the URI scheme does not support parallel streams or
             if a global override has been applied. If `False` parallel
             streams will be disabled.
+        tmpdir : `ResourcePath` or `None`, optional
+            Explicit override of the temporary directory to use for remote
+            downloads.
 
         Returns
         -------
@@ -1038,7 +1041,9 @@ def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
         raise NotImplementedError()
 
     @contextlib.contextmanager
-    def as_local(self, multithreaded: bool = True) -> Iterator[ResourcePath]:
+    def as_local(
+        self, multithreaded: bool = True, tmpdir: ResourcePathExpression | None = None
+    ) -> Iterator[ResourcePath]:
         """Return the location of the (possibly remote) resource as local file.
 
         Parameters
@@ -1049,6 +1054,9 @@ def as_local(self, multithreaded: bool = True) -> Iterator[ResourcePath]:
             effect if the URI scheme does not support parallel streams or
             if a global override has been applied. If `False` parallel
             streams will be disabled.
+        tmpdir : `ResourcePathExpression` or `None`, optional
+            Explicit override of the temporary directory to use for remote
+            downloads.
 
         Yields
         ------
@@ -1074,7 +1082,10 @@ def as_local(self, multithreaded: bool = True) -> Iterator[ResourcePath]:
         """
         if self.isdir():
             raise IsADirectoryError(f"Directory-like URI {self} cannot be fetched as local.")
-        local_src, is_temporary = self._as_local(multithreaded=multithreaded)
+        temp_dir = ResourcePath(tmpdir, forceDirectory=True) if tmpdir is not None else None
+        if temp_dir is not None and not temp_dir.isLocal:
+            raise ValueError(f"Temporary directory for as_local must be local resource not {temp_dir}")
+        local_src, is_temporary = self._as_local(multithreaded=multithreaded, tmpdir=temp_dir)
         local_uri = ResourcePath(local_src, isTemporary=is_temporary)
 
         try:
@@ -1100,8 +1111,9 @@ def temporary_uri(
         Parameters
         ----------
         prefix : `ResourcePath`, optional
-            Prefix to use. Without this the path will be formed as a local
-            file URI in a temporary directory. Ensuring that the prefix
+            Temporary directory to use (can be any scheme). Without this the
+            path will be formed as a local file URI in a temporary directory
+            created by `tempfile.mkdtemp`. Ensuring that the prefix
             location exists is the responsibility of the caller.
         suffix : `str`, optional
             A file suffix to be used. The ``.`` should be included in this

python/lsst/resources/file.py

@@ -79,7 +79,7 @@ def remove(self) -> None:
         """Remove the resource."""
         os.remove(self.ospath)
 
-    def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
+    def _as_local(self, multithreaded: bool = True, tmpdir: ResourcePath | None = None) -> tuple[str, bool]:
         """Return the local path of the file.
 
         This is an internal helper for ``as_local()``.
@@ -88,6 +88,8 @@ def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
         ----------
         multithreaded : `bool`, optional
             Unused.
+        tmpdir : `ResourcePath` or `None`, optional
+            Unused.
 
         Returns
         -------

python/lsst/resources/gs.py

@@ -202,9 +202,9 @@ def mkdir(self) -> None:
         # Should this method do anything at all?
         self.blob.upload_from_string(b"", retry=_RETRY_POLICY)
 
-    def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
+    def _as_local(self, multithreaded: bool = True, tmpdir: ResourcePath | None = None) -> tuple[str, bool]:
         with (
-            ResourcePath.temporary_uri(suffix=self.getExtension(), delete=False) as tmp_uri,
+            ResourcePath.temporary_uri(prefix=tmpdir, suffix=self.getExtension(), delete=False) as tmp_uri,
             time_this(log, msg="Downloading %s to local file", args=(self,)),
         ):
             try:

python/lsst/resources/http.py

@@ -98,6 +98,20 @@ def _timeout_from_environment(env_var: str, default_value: float) -> float:
     return timeout
 
 
+@functools.lru_cache
+def _calc_tmpdir_buffer_size(tmpdir: str) -> int:
+    """Compute the block size as 256 blocks of typical size
+    (i.e. 4096 bytes) or 10 times the file system block size,
+    whichever is higher.
+
+    This is a reasonable compromise between
+    using memory for buffering and the number of system calls
+    issued to read from or write to temporary files.
+    """
+    fsstats = os.statvfs(tmpdir)
+    return max(10 * fsstats.f_bsize, 256 * 4096)
+
+
 class HttpResourcePathConfig:
     """Configuration class to encapsulate the configurable items used by class
     HttpResourcePath.
@@ -381,8 +395,8 @@ def tmpdir_buffersize(self) -> tuple[str, int]:
         # whichever is higher. This is a reasonable compromise between
         # using memory for buffering and the number of system calls
         # issued to read from or write to temporary files.
-        fsstats = os.statvfs(tmpdir)
-        self._tmpdir_buffersize = (tmpdir, max(10 * fsstats.f_bsize, 256 * 4096))
+        bufsize = _calc_tmpdir_buffer_size(tmpdir)
+        self._tmpdir_buffersize = (tmpdir, bufsize)
 
         return self._tmpdir_buffersize
 
@@ -1480,17 +1494,20 @@ def _sign_with_macaroon(self, activity: ActivityCaveat, expiration_time_seconds:
         except json.JSONDecodeError:
             raise ValueError(f"could not deserialize response to POST request for URL {self}")
 
-    def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
+    def _as_local(self, multithreaded: bool = True, tmpdir: ResourcePath | None = None) -> tuple[str, bool]:
         """Download object over HTTP and place in temporary directory.
 
         Parameters
         ----------
         multithreaded : `bool`, optional
             If `True` the transfer will be allowed to attempt to improve
-          throughput by using parallel download streams. This may of no
+            throughput by using parallel download streams. This may of no
             effect if the URI scheme does not support parallel streams or
             if a global override has been applied. If `False` parallel
             streams will be disabled.
+        tmpdir : `ResourcePath` or `None`, optional
+            Explicit override of the temporary directory to use for remote
+            downloads.
 
         Returns
         -------
@@ -1509,9 +1526,14 @@ def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
                     f"Unable to download resource {self}; status: {resp.status_code} {resp.reason}"
                 )
 
-            tmpdir, buffer_size = self._config.tmpdir_buffersize
+            if tmpdir is None:
+                temp_dir, buffer_size = self._config.tmpdir_buffersize
+                tmpdir = ResourcePath(temp_dir, forceDirectory=True)
+            else:
+                buffer_size = _calc_tmpdir_buffer_size(tmpdir.ospath)
+
             with ResourcePath.temporary_uri(
-                suffix=self.getExtension(), prefix=ResourcePath(tmpdir, forceDirectory=True), delete=False
+                suffix=self.getExtension(), prefix=tmpdir, delete=False
             ) as tmp_uri:
                 expected_length = int(resp.headers.get("Content-Length", "-1"))
                 with time_this(

python/lsst/resources/mem.py

@@ -27,5 +27,5 @@ def exists(self) -> bool:
         """Test for existence and always return False."""
         return True
 
-    def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
+    def _as_local(self, multithreaded: bool = True, tmpdir: ResourcePath | None = None) -> tuple[str, bool]:
         raise RuntimeError(f"Do not know how to retrieve data for URI '{self}'")

python/lsst/resources/packageresource.py

@@ -29,7 +29,7 @@
         AbstractFileSystem = type
 
 from ._resourceHandles._baseResourceHandle import ResourceHandleProtocol
-from ._resourcePath import ResourcePath
+from ._resourcePath import ResourcePath, ResourcePathExpression
 
 log = logging.getLogger(__name__)
 
@@ -85,13 +85,17 @@ def read(self, size: int = -1) -> bytes:
             return fh.read(size)
 
     @contextlib.contextmanager
-    def as_local(self, multithreaded: bool = True) -> Iterator[ResourcePath]:
+    def as_local(
+        self, multithreaded: bool = True, tmpdir: ResourcePathExpression | None = None
+    ) -> Iterator[ResourcePath]:
         """Return the location of the Python resource as local file.
 
         Parameters
         ----------
         multithreaded : `bool`, optional
             Unused.
+        tmpdir : `ResourcePathExpression` or `None`, optional
+            Unused.
 
         Yields
         ------

python/lsst/resources/s3.py

@@ -353,7 +353,7 @@ def to_fsspec(self) -> tuple[AbstractFileSystem, str]:
 
         return s3, f"{self._bucket}/{self.relativeToPathRoot}"
 
-    def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
+    def _as_local(self, multithreaded: bool = True, tmpdir: ResourcePath | None = None) -> tuple[str, bool]:
         """Download object from S3 and place in temporary directory.
 
         Parameters
@@ -364,6 +364,9 @@ def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
             effect if the URI scheme does not support parallel streams or
             if a global override has been applied. If `False` parallel
             streams will be disabled.
+        tmpdir : `ResourcePath` or `None`, optional
+            Explicit override of the temporary directory to use for remote
+            downloads.
 
         Returns
         -------
@@ -373,7 +376,7 @@ def _as_local(self, multithreaded: bool = True) -> tuple[str, bool]:
             Always returns `True`. This is always a temporary file.
         """
         with (
-            ResourcePath.temporary_uri(suffix=self.getExtension(), delete=False) as tmp_uri,
+            ResourcePath.temporary_uri(prefix=tmpdir, suffix=self.getExtension(), delete=False) as tmp_uri,
             self._use_threads_temp_override(multithreaded),
             time_this(log, msg="Downloading %s to local file", args=(self,)),
         ):