added test to cleanup

flint/coadd/linmos.py

@@ -545,6 +545,26 @@ def generate_linmos_parameter_set(
     return linmos_parset_summary
 
 
+def _linmos_cleanup(linmos_parset_summary: LinmosParsetSummary) -> Tuple[Path, ...]:
+    """Clean up linmos files if requested.
+
+    Args:
+        linmos_parset_summary (LinmosParsetSummary): Parset summary from which the text file weights are gathered for deletion from
+
+    Returns:
+        Tuple[Path, ...]: Set of files removed
+    """
+
+    from flint.utils import remove_files_folders
+
+    removed_files = []
+    if linmos_parset_summary.weight_text_paths is not None:
+        removed_files.extend(
+            remove_files_folders(*linmos_parset_summary.weight_text_paths)
+        )
+    return tuple(removed_files)
+
+
 # TODO: These options are starting to get a little large. Perhaps we should use BaseOptions.
 def linmos_images(
     images: Collection[Path],
@@ -619,12 +639,7 @@ def linmos_images(
         )
 
     if cleanup:
-        if linmos_parset_summary.weight_text_paths is not None:
-            logger.info(
-                f"Remoing {len(linmos_parset_summary.weight_text_paths)} weight files generated"
-            )
-            for weight_file in linmos_parset_summary.weight_text_paths:
-                Path(weight_file).unlink()
+        _linmos_cleanup(linmos_parset_summary=linmos_parset_summary)
 
     return linmos_cmd
 

tests/test_linmos_coadd.py

@@ -11,6 +11,8 @@
 
 from flint.coadd.linmos import (
     BoundingBox,
+    LinmosParsetSummary,
+    _linmos_cleanup,
     _create_bound_box_plane,
     _get_alpha_linmos_option,
     _get_holography_linmos_options,
@@ -74,6 +76,30 @@ def test_linmos_alpha_option():
         _get_alpha_linmos_option(pol_axis=1234)
 
 
+def test_cleanup_image_weights_(tmpdir):
+    """Remove the weight files that have been created"""
+    cube_weight = Path(tmpdir) / "cubeweight"
+    cube_weight.mkdir(parents=True, exist_ok=True)
+    cube_fits = cube_weight / "cube.fits"
+
+    create_image_cube(out_path=cube_fits)
+    weight_file = cube_fits.with_suffix(".weights.txt")
+    assert not weight_file.exists()
+
+    weight_paths = generate_weights_list_and_files(image_paths=[cube_fits], mode="mad")
+    linmos_parset_summary = LinmosParsetSummary(
+        parset_path=Path("JackSparrow.txt"),
+        image_paths=tuple([cube_fits]),
+        weight_text_paths=weight_paths,
+    )
+    assert weight_file.exists()
+    assert isinstance(linmos_parset_summary, LinmosParsetSummary)
+    files_removed = _linmos_cleanup(linmos_parset_summary=linmos_parset_summary)
+    assert not weight_file.exists()
+    assert len(files_removed) == 1
+    assert files_removed[0] == weight_file
+
+
 def test_get_image_weights(tmpdir):
     """See whether the weights computed per plane in a cube work appropriately"""
     cube_weight = Path(tmpdir) / "cubeweight"