From 50a14561650e83fe6681f9ac26087f21396b17c7 Mon Sep 17 00:00:00 2001 From: Leonardo Schwarz Date: Mon, 1 Jul 2024 16:49:22 +0200 Subject: [PATCH] code style: parallel_ops --- .../parallel_ops/read_spectra_parallel.py | 8 ++-- .../parallel_ops/write_spectra_parallel.py | 47 ++++++++++--------- src/depiction/tools/cli/correct_baseline.py | 6 ++- src/depiction/tools/correct_baseline.py | 12 ++--- src/depiction/tools/create_imzml_pool.py | 2 +- src/depiction/tools/split_imzml.py | 1 + 6 files changed, 43 insertions(+), 33 deletions(-) diff --git a/src/depiction/parallel_ops/read_spectra_parallel.py b/src/depiction/parallel_ops/read_spectra_parallel.py index 29003c5..ab49e7f 100644 --- a/src/depiction/parallel_ops/read_spectra_parallel.py +++ b/src/depiction/parallel_ops/read_spectra_parallel.py @@ -1,9 +1,11 @@ from __future__ import annotations + from typing import ( Any, Callable, TypeVar, TYPE_CHECKING, + TypedDict, ) import numpy as np @@ -28,11 +30,11 @@ def __init__(self, config: ParallelConfig) -> None: self._config = config @classmethod - def from_config(cls, config: ParallelConfig): + def from_config(cls, config: ParallelConfig) -> ReadSpectraParallel: return cls(config=config) @classmethod - def from_params(cls, n_jobs: int, task_size: int | None, verbose: int = 1): + def from_params(cls, n_jobs: int, task_size: int | None, verbose: int = 1) -> ReadSpectraParallel: """In general, try to use from_config and pass the configuration throughout the application as appropriate.""" return cls(config=ParallelConfig(n_jobs=n_jobs, task_size=task_size, verbose=verbose)) @@ -76,7 +78,7 @@ def map_chunked( for task_index, task in enumerate(self._config.get_task_splits(item_indices=spectra_indices)) ] - def execute_task(args, **kwargs) -> list[T]: + def execute_task(args: list[Any], **kwargs: TypedDict[str, Any]) -> list[T]: with read_file.reader() as reader: return operation(reader, *args, **kwargs) diff --git a/src/depiction/parallel_ops/write_spectra_parallel.py b/src/depiction/parallel_ops/write_spectra_parallel.py index 1c24d2c..c927dae 100644 --- a/src/depiction/parallel_ops/write_spectra_parallel.py +++ b/src/depiction/parallel_ops/write_spectra_parallel.py @@ -1,12 +1,11 @@ +from __future__ import annotations import contextlib import functools import os -from typing import Optional, Callable, Any, Union +from typing import Callable, Any, TYPE_CHECKING from tempfile import TemporaryDirectory -from numpy.typing import NDArray from depiction.parallel_ops import ReadSpectraParallel -from depiction.parallel_ops.parallel_config import ParallelConfig from depiction.persistence import ( ImzmlReadFile, ImzmlWriteFile, @@ -16,25 +15,29 @@ ) from depiction.tools.merge_imzml import MergeImzml +if TYPE_CHECKING: + from numpy.typing import NDArray + from depiction.parallel_ops.parallel_config import ParallelConfig + class WriteSpectraParallel: def __init__(self, config: ParallelConfig) -> None: self._config = config @classmethod - def from_config(cls, config: ParallelConfig): + def from_config(cls, config: ParallelConfig) -> WriteSpectraParallel: return cls(config) def map_chunked_to_files( self, read_file: ImzmlReadFile, write_files: list[ImzmlWriteFile], - operation: Union[ - Callable[[ImzmlReader, list[int], list[ImzmlWriter], ...], None], - Callable[[ImzmlReader, list[int], list[ImzmlWriteFile], ...], None], - ], - spectra_indices: Optional[NDArray[int]] = None, - bind_args: Optional[dict[str, Any]] = None, + operation: ( + Callable[[ImzmlReader, list[int], list[ImzmlWriter], ...], None] + | Callable[[ImzmlReader, list[int], list[ImzmlWriteFile], ...], None] + ), + spectra_indices: NDArray[int] | None = None, + bind_args: dict[str, Any] | None = None, open_write_files: bool = True, ) -> None: """Maps an operation over a file, in chunks, writing the results to a list of files. @@ -81,14 +84,14 @@ def map_chunked_external_to_files( read_file: ImzmlReadFile, write_files: list[ImzmlWriteFile], operation: Callable[[str, list[str]], None], - spectra_indices: Optional[NDArray[int]] = None, - bind_args: Optional[dict[str, Any]] = None, + spectra_indices: NDArray[int] | None = None, + bind_args: dict[str, Any] | None = None, ) -> None: def op( reader: ImzmlReader, spectra_ids: list[int], write_files: list[ImzmlWriteFile], - **kwargs, + **kwargs: dict[str, Any], ) -> None: # TODO maybe kwarg handling could be done a bit more clean here in the future # TODO also it's currently untested @@ -120,7 +123,7 @@ def _get_split_modes_and_paths( work_directory: str, read_file: ImzmlReadFile, write_files: list[ImzmlWriteFile], - spectra_indices: Optional[NDArray[int]], + spectra_indices: NDArray[int] | None, ) -> list[tuple[ImzmlModeEnum, list[str]]]: # determine the number of tasks if spectra_indices is not None: @@ -142,10 +145,10 @@ def _write_transformed_chunked_operation( reader: ImzmlReader, spectra_indices: list[int], task_index: int, - operation: Union[ - Callable[[ImzmlReader, list[int], list[ImzmlWriter], ...], None], - Callable[[ImzmlReader, list[int], list[ImzmlWriteFile], ...], None], - ], + operation: ( + Callable[[ImzmlReader, list[int], list[ImzmlWriter], ...], None] + | Callable[[ImzmlReader, list[int], list[ImzmlWriteFile], ...], None] + ), open_write_files: bool, split_modes_and_paths: list[tuple[ImzmlModeEnum, list[str]]], ) -> None: @@ -192,10 +195,12 @@ def map_chunked_to_file( read_file: ImzmlReadFile, write_file: ImzmlWriteFile, operation: Callable[[ImzmlReader, list[int], ImzmlWriter], None], - spectra_indices: Optional[NDArray[int]] = None, - bind_args: Optional[dict[str, Any]] = None, + spectra_indices: NDArray[int] | None = None, + bind_args: dict[str, Any] | None = None, ): - def wrap_operation(reader: ImzmlReader, spectra_ids: list[int], writers: list[ImzmlWriter], **kwargs): + def wrap_operation( + reader: ImzmlReader, spectra_ids: list[int], writers: list[ImzmlWriter], **kwargs: dict[str, Any] + ): return operation(reader, spectra_ids, writers[0], **kwargs) return self.map_chunked_to_files( diff --git a/src/depiction/tools/cli/correct_baseline.py b/src/depiction/tools/cli/correct_baseline.py index 64cf093..834e9d9 100644 --- a/src/depiction/tools/cli/correct_baseline.py +++ b/src/depiction/tools/cli/correct_baseline.py @@ -1,8 +1,7 @@ from __future__ import annotations import shutil -from pathlib import Path -from typing import Annotated, Literal +from typing import Annotated, Literal, TYPE_CHECKING import typer from loguru import logger @@ -12,6 +11,9 @@ from depiction.persistence import ImzmlReadFile, ImzmlWriteFile from depiction.tools.correct_baseline import BaselineVariants, CorrectBaseline +if TYPE_CHECKING: + from pathlib import Path + def correct_baseline( input_imzml: Annotated[Path, Argument()], diff --git a/src/depiction/tools/correct_baseline.py b/src/depiction/tools/correct_baseline.py index 41f587e..4cdf606 100644 --- a/src/depiction/tools/correct_baseline.py +++ b/src/depiction/tools/correct_baseline.py @@ -7,16 +7,16 @@ from depiction.parallel_ops.parallel_config import ParallelConfig from depiction.parallel_ops.write_spectra_parallel import WriteSpectraParallel -from depiction.persistence import ( - ImzmlReadFile, - ImzmlWriteFile, - ImzmlWriter, - ImzmlReader, -) from depiction.spectrum.baseline.local_medians_baseline import LocalMediansBaseline from depiction.spectrum.baseline.tophat_baseline import TophatBaseline if TYPE_CHECKING: + from depiction.persistence import ( + ImzmlReadFile, + ImzmlWriteFile, + ImzmlWriter, + ImzmlReader, + ) from numpy.typing import NDArray from depiction.spectrum.baseline.baseline import Baseline diff --git a/src/depiction/tools/create_imzml_pool.py b/src/depiction/tools/create_imzml_pool.py index f99ba10..9a7583b 100644 --- a/src/depiction/tools/create_imzml_pool.py +++ b/src/depiction/tools/create_imzml_pool.py @@ -64,7 +64,7 @@ def _write_imzml_file(self, write_file: ImzmlWriteFile) -> None: with write_file.writer() as writer: for imzml_file in tqdm(self._source_files, desc="Copying file"): with imzml_file.reader() as reader: - abs_path = str(imzml_file.imzml_file.absolute()) + str(imzml_file.imzml_file.absolute()) spectrum_ids = self.pool_source_df.query("abs_path == @abs_path").iloc[0]["source_spectrum_id"] writer.copy_spectra(reader, spectrum_ids) diff --git a/src/depiction/tools/split_imzml.py b/src/depiction/tools/split_imzml.py index d59e02e..24202d2 100644 --- a/src/depiction/tools/split_imzml.py +++ b/src/depiction/tools/split_imzml.py @@ -2,6 +2,7 @@ import logging import os +from pathlib import Path from typing import Optional import numpy as np