Merge pull request #129 from adl1995/io/overwrite

Add overwrite option in write_image function for FITS files

hips/draw/tests/test_ui.py

@@ -71,6 +71,11 @@ def test_make_sky_image(tmpdir, pars):
     assert_allclose(np.sum(result.image, dtype=float), pars['data_sum'])
     assert_allclose(result.image[200, 994], pars['data_1'])
     assert_allclose(result.image[200, 995], pars['data_2'])
-    result.write_image(str(tmpdir / 'test.' + pars['file_format']))
+    result.write_image(str(tmpdir / 'test.' + pars['file_format']), False)
+
+    # Try and overwrite the image.
+    with pytest.raises(FileExistsError):
+        result.write_image(str(tmpdir / 'test.' + pars['file_format']), False)
+
     result.plot()
     result.report()

hips/draw/ui.py

@@ -2,6 +2,7 @@
 """The high-level end user interface (UI)."""
 import numpy as np
 from PIL import Image
+from pathlib import Path
 from astropy.io import fits
 from typing import List, Union
 from ..utils.wcs import WCSGeometry
@@ -100,14 +101,19 @@ def from_painter(cls, painter: HipsPainter) -> 'HipsDrawResult':
             stats=painter._stats,
         )
 
-    def write_image(self, filename: str) -> None:
+    def write_image(self, filename: str, overwrite: bool = False) -> None:
         """Write image to file.
 
         Parameters
         ----------
         filename : str
             Filename
+        overwrite : bool
+            Overwrite the output file, if it exists
         """
+        if overwrite == False and Path(filename).exists():
+            raise FileExistsError(f"File {filename} already exists.")
+
         if self.tile_format == 'fits':
             hdu = fits.PrimaryHDU(data=self.image, header=self.geometry.fits_header)
             hdu.writeto(filename)