support overwriting work_unit shards (#783)

dirac-institute · Jan 24, 2025 · 2c8370d · 2c8370d
1 parent 48aa958
commit 2c8370d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/kbmod/work_unit.py b/src/kbmod/work_unit.py
@@ -584,7 +584,7 @@ def to_sharded_fits(self, filename, directory, overwrite=False):
             psf_hdu = fits.hdu.image.ImageHDU(psf_array)
             psf_hdu.name = f"PSF_{i}"
             sub_hdul.append(psf_hdu)
-            sub_hdul.writeto(os.path.join(directory, f"{i}_{filename}"))
+            sub_hdul.writeto(os.path.join(directory, f"{i}_{filename}"), overwrite=overwrite)
 
         # Create a primary file with all of the metadata, including all the WCS info.
         hdul = self.metadata_to_hdul()

diff --git a/tests/test_work_unit.py b/tests/test_work_unit.py
@@ -336,7 +336,7 @@ def test_save_and_load_fits_shard(self):
             self.assertRaises(FileExistsError, work.to_fits, file_path)
 
             # We succeed if overwrite=True
-            work.to_fits(file_path, overwrite=True)
+            work.to_sharded_fits("test_workunit.fits", dir_name, overwrite=True)
 
     def test_save_and_load_fits_shard_lazy(self):
         with tempfile.TemporaryDirectory() as dir_name: