diff --git a/pygmt/figure.py b/pygmt/figure.py index 09b5b4645d1..ade6a31e35e 100644 --- a/pygmt/figure.py +++ b/pygmt/figure.py @@ -252,7 +252,14 @@ def psconvert(self, **kwargs): ) def savefig( - self, fname, transparent=False, crop=True, anti_alias=True, show=False, **kwargs + self, + fname, + transparent=False, + crop=True, + anti_alias=True, + show=False, + worldfile=False, + **kwargs, ): """ Save the figure to an image file. @@ -297,6 +304,14 @@ def savefig( :meth:`pygmt.Figure.psconvert`. Ignored if creating vector images. show: bool If ``True``, will open the figure in an external viewer. + worldfile : bool + If ``True``, will create a companion + `world file `__ for the + figure. The world file will have the same name as the figure file + but with different extension (e.g. tfw for tif). See + https://en.wikipedia.org/wiki/World_file#Filename_extension + for the convention of world file extensions. This parameter only + works for raster image formats (except GeoTIFF). dpi : int Set raster resolution in dpi [Default is ``720`` for PDF, ``300`` for others]. @@ -305,6 +320,7 @@ def savefig( :meth:`pygmt.Figure.psconvert`. Valid parameters are ``gs_path``, ``gs_option``, ``resize``, ``bb_style``, and ``verbose``. """ + # pylint: disable=too-many-branches # All supported formats fmts = { "bmp": "b", @@ -347,6 +363,13 @@ def savefig( kwargs["Qt"] = 2 kwargs["Qg"] = 2 + if worldfile: + if ext in ["eps", "kml", "pdf", "tiff"]: + raise GMTInvalidInput( + f"Saving a world file is not supported for '{ext}' format." + ) + kwargs["W"] = True + self.psconvert(prefix=prefix, fmt=fmt, crop=crop, **kwargs) # Remove the .pgw world file if exists diff --git a/pygmt/tests/test_figure.py b/pygmt/tests/test_figure.py index 95d2241fa35..93e044f9056 100644 --- a/pygmt/tests/test_figure.py +++ b/pygmt/tests/test_figure.py @@ -292,6 +292,27 @@ def mock_psconvert(*args, **kwargs): # pylint: disable=unused-argument } +def test_figure_savefig_worldfile(): + """ + Check if a world file is created for supported formats and raise an error + for unsupported formats. + """ + fig = Figure() + fig.basemap(region=[0, 1, 0, 1], projection="X1c/1c", frame=True) + # supported formats + for fmt in [".bmp", ".jpg", ".png", ".ppm", ".tif"]: + with GMTTempFile(prefix="pygmt-worldfile", suffix=fmt) as imgfile: + fig.savefig(fname=imgfile.name, worldfile=True) + assert Path(imgfile.name).stat().st_size > 0 + worldfile_suffix = "." + fmt[1] + fmt[3] + "w" + assert Path(imgfile.name).with_suffix(worldfile_suffix).stat().st_size > 0 + # unsupported formats + for fmt in [".eps", ".kml", ".pdf", ".tiff"]: + with GMTTempFile(prefix="pygmt-worldfile", suffix=fmt) as imgfile: + with pytest.raises(GMTInvalidInput): + fig.savefig(fname=imgfile.name, worldfile=True) + + @pytest.mark.skipif(IPython is None, reason="run when IPython is installed") def test_figure_show(): """