Figure.savefig: Add the 'worldfile' parameter to write a companion wo…

…rld file for raster images (#2766) Co-authored-by: Yvonne Fröhlich <[email protected]> Co-authored-by: Wei Ji <[email protected]>
GenericMappingTools · Oct 31, 2023 · a9bd688 · a9bd688
1 parent 3d2fd89
commit a9bd688
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 1 deletion.
diff --git 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 <https://en.wikipedia.org/wiki/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
@@ -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():
     """