diff --git a/hips/draw/tests/test_ui.py b/hips/draw/tests/test_ui.py index a701354..786f3d2 100644 --- a/hips/draw/tests/test_ui.py +++ b/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() diff --git a/hips/draw/ui.py b/hips/draw/ui.py index 69dd9ff..50413d6 100644 --- a/hips/draw/ui.py +++ b/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)