#146 Merge pull request from deshima-dev/astropenguin/issue131

Add overwrite option to demerge function

CITATION.cff

@@ -31,5 +31,5 @@ keywords:
   - spectroscopy
   - deshima
 license: MIT
-version: 3.0.5
-date-released: '2024-07-11'
+version: 3.1.0
+date-released: '2024-07-14'

README.md

@@ -11,7 +11,7 @@ DESHIMA merge code for observed datasets
 ## Installation
 
 ```shell
-pip install demerge==3.0.5
+pip install demerge==3.1.0
 ```
 
 ## Command line interface

demerge/__init__.py

@@ -1,8 +1,9 @@
 __all__ = ["demerge", "merge", "reduce"]
-__version__ = "3.0.5"
+__version__ = "3.1.0"
 
 
 # standard library
+from contextlib import contextmanager
 from logging import DEBUG, basicConfig, getLogger
 from pathlib import Path
 from typing import Any
@@ -19,6 +20,19 @@
 LOGGER = getLogger(__name__)
 
 
+@contextmanager
+def set_logger(debug: bool):
+    level = LOGGER.level
+
+    if debug:
+        LOGGER.setLevel(DEBUG)
+
+    try:
+        yield
+    finally:
+        LOGGER.setLevel(level)
+
+
 def demerge(
     obsid: str,
     /,
@@ -27,6 +41,7 @@ def demerge(
     dems_dir: Path = Path(),
     reduced_dir: Path = Path(),
     ddb: Path = DEFAULT_DDB_FILE,
+    overwrite: bool = False,
     debug: bool = False,
     **merge_options: Any,
 ) -> Path:
@@ -41,31 +56,31 @@ def demerge(
         reduced_dir: Path where reduced data directory will be placed,
             i.e. expecting ``${reduced_dir}/reduced_YYYYmmddHHMMSS``.
         ddb: Path of DDB (DESHIMA database) file.
+        overwrite: If True, reduced data directory and merged DEMS file
+            will be overwritten even if they exist.
         debug: If True, detailed logs for debugging will be printed.
         **merge_options: Other merge options for the merge command.
 
     Returns:
         Path of the merged DEMS file.
 
     """
-    if debug:
-        LOGGER.setLevel(DEBUG)
-
-    basicConfig(
-        datefmt="%Y-%m-%d %H:%M:%S",
-        format="[%(asctime)s %(name)s %(levelname)s] %(message)s",
-    )
-
-    for key, val in locals().items():
-        LOGGER.debug(f"{key}: {val!r}")
+    with set_logger(debug):
+        for key, val in locals().items():
+            LOGGER.debug(f"{key}: {val!r}")
 
     data_dir_ = Path(data_dir).resolve() / f"cosmos_{obsid}"
     reduced_dir_ = Path(reduced_dir).resolve() / f"reduced_{obsid}"
     dems_dir_ = Path(dems_dir).resolve()
     ddb = Path(ddb).resolve()
 
     # Run reduce function
-    readout = reduce.reduce(data_dir_, reduced_dir_, debug=debug)
+    readout = reduce.reduce(
+        data_dir_,
+        reduced_dir_,
+        overwrite=overwrite,
+        debug=debug,
+    )
 
     # Run merge function
     if (dems := dems_dir_ / f"dems_{obsid}.zarr.zip").exists():
@@ -111,11 +126,17 @@ def demerge(
         weather=weather,
         misti=misti,
         cabin=cabin,
+        overwrite=overwrite,
         debug=debug,
         **merge_options,
     )
 
 
 def cli() -> None:
     """Command line interface of the demerge function."""
+    basicConfig(
+        datefmt="%Y-%m-%d %H:%M:%S",
+        format="[%(asctime)s %(name)s %(funcName)s %(levelname)s] %(message)s",
+    )
+
     Fire(demerge)

demerge/merge/__init__.py

@@ -2,6 +2,7 @@
 
 
 # standard library
+from contextlib import contextmanager
 from logging import DEBUG, basicConfig, getLogger
 from pathlib import Path
 
@@ -15,6 +16,19 @@
 LOGGER = getLogger(__name__)
 
 
+@contextmanager
+def set_logger(debug: bool):
+    level = LOGGER.level
+
+    if debug:
+        LOGGER.setLevel(DEBUG)
+
+    try:
+        yield
+    finally:
+        LOGGER.setLevel(level)
+
+
 def merge(
     dems: Path,
     /,
@@ -30,6 +44,7 @@ def merge(
     cabin: str = "",
     coordinate: str = "azel",
     measure: str = "df/f",
+    overwrite: bool = False,
     debug: bool = False,
     offset_time_antenna: int = 0,
 ) -> Path:
@@ -48,6 +63,7 @@ def merge(
         cabin: Path of the cabin log file (.cabin).
         coordinate: Coordinate system of the output data (azel or radec).
         measure: Output data format (df/f or brightness).
+        overwrite: If True, ``dems`` will be overwritten even if it exists.
         debug: If True, detailed logs for debugging will be printed.
         offset_time_antenna: Time diff (ms) between reduced FITS and antenna log.
 
@@ -58,20 +74,10 @@ def merge(
         FileExistsError: Raised if ``dems`` exists.
 
     """
-    # ロガーの設定
-    if debug:
-        LOGGER.setLevel(DEBUG)
-
-    basicConfig(
-        datefmt="%Y-%m-%d %H:%M:%S",
-        format="[%(asctime)s %(name)s %(levelname)s] %(message)s",
-    )
+    with set_logger(debug):
+        for key, val in locals().items():
+            LOGGER.debug(f"{key}: {val!r}")
 
-    # 引数と値をロガーに記録
-    for key, val in locals().items():
-        LOGGER.debug(f"{key}: {val!r}")
-
-    # マージの実行
     da = create_dems(
         ddbfits_path=ddb,
         corresp_path=corresp,
@@ -87,14 +93,22 @@ def merge(
         offset_time_antenna=offset_time_antenna,
     )
 
-    if (dems := Path(dems)).exists():
+    if (dems := Path(dems)).exists() and not overwrite:
         raise FileExistsError(dems)
 
+    if overwrite:
+        dems.unlink(missing_ok=True)
+
     dems.parent.mkdir(exist_ok=True, parents=True)
     da.to_zarr(dems, mode="w")
     return dems.resolve()
 
 
 def cli() -> None:
     """Command line interface of the merge function."""
+    basicConfig(
+        datefmt="%Y-%m-%d %H:%M:%S",
+        format="[%(asctime)s %(name)s %(funcName)s %(levelname)s] %(message)s",
+    )
+
     Fire(merge)

demerge/reduce/__init__.py

@@ -2,8 +2,10 @@
 
 
 # standard library
+from contextlib import contextmanager
 from logging import DEBUG, basicConfig, getLogger
 from pathlib import Path
+from shutil import rmtree
 from subprocess import run
 from tempfile import TemporaryDirectory
 
@@ -17,18 +19,33 @@
 SCRIPTS = Path(__file__).parent / "utils" / "scripts" / "aste"
 
 
+@contextmanager
+def set_logger(debug: bool):
+    level = LOGGER.level
+
+    if debug:
+        LOGGER.setLevel(DEBUG)
+
+    try:
+        yield
+    finally:
+        LOGGER.setLevel(level)
+
+
 def reduce(
     data_dir: Path,
     reduced_dir: Path,
     /,
     *,
+    overwrite: bool = False,
     debug: bool = False,
 ) -> Path:
     """Reduce raw data of KID measurements into a single "reduced" FITS.
 
     Args:
         data_dir: Path of raw data directory (e.g. ``cosmos_YYYYmmddHHMMSS``).
         reduced_dir: Path of reduced data directory (e.g. ``reduced_YYYYmmddHHMMSS``).
+        overwrite: If True, ``reduced_dir`` will be overwritten even if it exists.
         debug: If True, detailed logs for debugging will be printed.
 
     Returns:
@@ -39,24 +56,20 @@ def reduce(
         FileExistsError: Raised if ``reduced_dir`` exists.
 
     """
-    if debug:
-        LOGGER.setLevel(DEBUG)
-
-    basicConfig(
-        datefmt="%Y-%m-%d %H:%M:%S",
-        format="[%(asctime)s %(name)s %(levelname)s] %(message)s",
-    )
-
-    for key, val in locals().items():
-        LOGGER.debug(f"{key}: {val!r}")
+    with set_logger(debug):
+        for key, val in locals().items():
+            LOGGER.debug(f"{key}: {val!r}")
 
     # Resolve paths (must be done before changing working directory)
     if not (data_dir := Path(data_dir).resolve()).exists():
         raise FileNotFoundError(data_dir)
 
-    if (reduced_dir := Path(reduced_dir).resolve()).exists():
+    if (reduced_dir := Path(reduced_dir).resolve()).exists() and not overwrite:
         raise FileExistsError(reduced_dir)
 
+    if overwrite:
+        rmtree(reduced_dir, ignore_errors=True)
+
     # Run scripts in a temporary directory (to isolate intermediate files)
     with TemporaryDirectory() as work_dir:
         run(
@@ -86,4 +99,9 @@ def reduce(
 
 def cli() -> None:
     """Command line interface of the reduce function."""
+    basicConfig(
+        datefmt="%Y-%m-%d %H:%M:%S",
+        format="[%(asctime)s %(name)s %(funcName)s %(levelname)s] %(message)s",
+    )
+
     Fire(reduce)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "demerge"
-version = "3.0.5"
+version = "3.1.0"
 description = "DESHIMA merge code for observed datasets"
 authors = [
     "Tatsuya Takekoshi <tatsuya.takekoshi@gmail.com>",