Add processing directory argument

src/sim_recon/cli/parsing/recon.py

@@ -60,6 +60,11 @@ def parse_args(
         "--output-directory",
         help="If specified, the output directory in which the reconstructed files will be saved (otherwise each reconstruction will be saved in the same directory as its SIM data file)",
     )
+    parser.add_argument(
+        "-p",
+        "--processing-directory",
+        help="If specified, the directory in which the temporary files will be stored for processing (otherwise the output directory will be used)",
+    )
     parser.add_argument(
         "--otf",
         dest="otfs",

src/sim_recon/cli/recon.py

@@ -16,6 +16,7 @@ def main() -> None:
         *namespace.sim_data_paths,
         config_path=namespace.config_path,
         output_directory=namespace.output_directory,
+        processing_directory=namespace.processing_directory,
         otf_overrides={} if namespace.otfs is None else dict(namespace.otfs),
         overwrite=namespace.overwrite,
         cleanup=namespace.cleanup,

src/sim_recon/files/utils.py

@@ -8,10 +8,13 @@
 from pathlib import Path
 from uuid import uuid4
 from contextlib import contextmanager
+from tempfile import TemporaryDirectory
+import weakref as _weakref
 from typing import TYPE_CHECKING
 
 from ..exceptions import (
     PySimReconFileExistsError,
+    PySimReconFileNotFoundError,
     PySimReconIOError,
     PySimReconOSError,
     PySimReconValueError,
@@ -184,3 +187,73 @@ def combine_text_files(
             f.write(header + separator)
         f.write(separator.join(contents_generator))
         os.fsync(f.fileno())
+
+
+class NamedTemporaryDirectory(TemporaryDirectory):
+
+    def __init__(
+        self,
+        directory: str | PathLike[str],
+        name: str | None = None,
+        parents: bool = True,
+        allow_fallback: bool = True,
+        ignore_cleanup_errors: bool = False,
+        *,
+        delete: bool = True,
+    ) -> None:
+        self.name: str
+        if name is not None:
+            path = Path(directory) / name
+            if not path.parent.is_dir():
+                if not parents:
+                    raise PySimReconFileNotFoundError(
+                        f"Parent directory {path.parent} does not exist"
+                    )
+                path.parent.mkdir(parents=True)
+            if not path.exists():
+                path.mkdir(exist_ok=False, parents=False)
+                self.name = str(path)
+                self._ignore_cleanup_errors = ignore_cleanup_errors
+                self._delete = delete
+                self._finalizer = _weakref.finalize(
+                    self,
+                    self._cleanup,  # type: ignore
+                    self.name,
+                    warn_message="Implicitly cleaning up {!r}".format(self),
+                    ignore_errors=self._ignore_cleanup_errors,
+                    delete=self._delete,
+                )
+                return
+            elif not allow_fallback:
+                raise PySimReconFileExistsError(
+                    f"Directory cannot be created as '{path}' exists"
+                )
+        # Used if name is not defined or if the path already exists and allow_fallback is True
+        super().__init__(
+            prefix=name,
+            dir=directory,
+            ignore_cleanup_errors=ignore_cleanup_errors,
+            delete=delete,
+        )
+
+
+@contextmanager
+def delete_directory_if_empty(
+    path: str | PathLike[str] | None,
+) -> Generator[None, None, None]:
+    if path is None:
+        yield None
+        return
+    try_cleanup = not os.path.isdir(path)
+    try:
+        yield None
+    finally:
+        if try_cleanup and os.path.isdir(path):
+            with os.scandir(path) as it:
+                directory_empty = not any(it)  # False if any entries found
+            if directory_empty:
+                logger.info(
+                    "Removing empty directory '%s'",
+                    path,
+                )
+                os.rmdir(path)

src/sim_recon/main.py

@@ -96,6 +96,7 @@ def sim_reconstruct(
     *sim_data_paths: str | PathLike[str],
     config_path: str | PathLike[str] | None = None,
     output_directory: str | PathLike[str] | None = None,
+    processing_directory: str | PathLike[str] | None = None,
     otf_overrides: dict[int, Path] | None = None,
     overwrite: bool = False,
     cleanup: bool = True,
@@ -116,6 +117,8 @@ def sim_reconstruct(
         Path of the top level config file, by default None
     output_directory : str | PathLike[str] | None, optional
         Directory to save reconstructions in (reconstructions will be saved with the data files if not specified), by default None
+    processing_directory : str | PathLike[str] | None, optional
+        The directory in which the temporary files will be stored for processing (otherwise the output directory will be used), by default None
     otf_overrides : dict[int, Path] | None, optional
         A dictionary with emission wavelengths in nm as keys and paths to OTF files as values (these override configured OTFs), by default None
     overwrite : bool, optional
@@ -137,6 +140,7 @@ def sim_reconstruct(
         conf,
         *sim_data_paths,
         output_directory=output_directory,
+        processing_directory=processing_directory,
         overwrite=overwrite,
         cleanup=cleanup,
         stitch_channels=stitch_channels,

src/sim_recon/recon.py

@@ -9,12 +9,17 @@
 from shutil import copyfile
 from pathlib import Path
 import numpy as np
-from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING
 
 from pycudasirecon.sim_reconstructor import SIMReconstructor, lib  # type: ignore[import-untyped]
 
-from .files.utils import redirect_output_to, create_output_path, combine_text_files
+from .files.utils import (
+    redirect_output_to,
+    create_output_path,
+    combine_text_files,
+    delete_directory_if_empty,
+    NamedTemporaryDirectory,
+)
 from .files.config import create_wavelength_config
 from .images import get_image_data, dv_to_tiff
 from .images.dv import write_dv, image_resolution_from_mrc, read_mrc_bound_array
@@ -363,6 +368,7 @@ def run_reconstructions(
     conf: ConfigManager,
     *sim_data_paths: str | PathLike[str],
     output_directory: str | PathLike[str] | None,
+    processing_directory: str | PathLike[str] | None = None,
     overwrite: bool = False,
     cleanup: bool = False,
     stitch_channels: bool = True,
@@ -384,6 +390,7 @@ def run_reconstructions(
             maxtasksperchild=1,
         ) as pool,
         logging_redirect(),
+        delete_directory_if_empty(processing_directory),
     ):
         for sim_data_path in progress_wrapper(
             sim_data_paths, desc="SIM data files", unit="file"
@@ -400,22 +407,23 @@ def run_reconstructions(
                     raise PySimReconFileNotFoundError(
                         f"Image file {sim_data_path} does not exist"
                     )
-
-                processing_directory: str | Path
-                with TemporaryDirectory(
-                    prefix="proc_",
-                    suffix=f"_{sim_data_path.stem}",
-                    dir=file_output_directory,
+                with NamedTemporaryDirectory(
+                    name=sim_data_path.stem,
+                    parents=True,
+                    directory=(
+                        file_output_directory
+                        if processing_directory is None
+                        else processing_directory
+                    ),
                     delete=cleanup,
-                ) as processing_directory:
-                    processing_directory = Path(processing_directory)
+                ) as proc_dir:
+                    proc_dir = Path(proc_dir)
 
-                    # These processing files are cleaned up by TemporaryDirectory
-                    # As single-wavelength files will be used directly and we don't
-                    # want to delete real input files!
+                    # These processing files are cleaned up by
+                    # NamedTemporaryDirectory if cleanup == True
                     processing_info_dict = _prepare_files(
                         sim_data_path,
-                        processing_directory,
+                        proc_dir,
                         conf=conf,
                         allow_missing_channels=allow_missing_channels,
                         **config_kwargs,
@@ -485,7 +493,7 @@ def run_reconstructions(
                         if not proc_log_files:
                             logger.warning(
                                 "No output log file created as no per-channel log files were found",
-                                processing_directory,
+                                proc_dir,
                             )
                         else:
                             log_path = create_output_path(
@@ -509,20 +517,6 @@ def run_reconstructions(
                                 "Reconstruction log file created at '%s'", log_path
                             )
 
-                    if cleanup:
-                        for processing_info in processing_info_dict.values():
-                            try:
-                                if processing_info.output_path.is_file():
-                                    logger.debug(
-                                        "Removing %s", processing_info.output_path
-                                    )
-                                    os.remove(processing_info.output_path)
-                            except Exception:
-                                logger.error(
-                                    "Failed to remove %s",
-                                    processing_info.output_path,
-                                    exc_info=True,
-                                )
             except ConfigException as e:
                 logger.error("Unable to process %s: %s", sim_data_path, e)
             except PySimReconException as e: