diff --git a/src/sim_recon/cli/parsing/recon.py b/src/sim_recon/cli/parsing/recon.py index 8ab40ca..9cda924 100644 --- a/src/sim_recon/cli/parsing/recon.py +++ b/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", diff --git a/src/sim_recon/cli/recon.py b/src/sim_recon/cli/recon.py index 4ee3277..216771a 100644 --- a/src/sim_recon/cli/recon.py +++ b/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, diff --git a/src/sim_recon/files/utils.py b/src/sim_recon/files/utils.py index 32de187..a659201 100644 --- a/src/sim_recon/files/utils.py +++ b/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) diff --git a/src/sim_recon/main.py b/src/sim_recon/main.py index 8615111..7f3c2cf 100644 --- a/src/sim_recon/main.py +++ b/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, diff --git a/src/sim_recon/recon.py b/src/sim_recon/recon.py index 0ee42fd..0a9ea6b 100644 --- a/src/sim_recon/recon.py +++ b/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: