Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Figure.savefig: Add the 'worldfile' parameter to write a companion world file for raster images #2766

Merged
merged 15 commits into from
Oct 31, 2023
Merged
25 changes: 24 additions & 1 deletion pygmt/figure.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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].
Expand All @@ -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",
Expand Down Expand Up @@ -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
Expand Down
21 changes: 21 additions & 0 deletions pygmt/tests/test_figure.py
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,27 @@ def mock_psconvert(*args, **kwargs): # pylint: disable=unused-argument
}


def test_figure_savefig_worldfile():
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was there a reason to switch from the parametrized tests to a single test in a863132? It'll be harder to see which formats/extensions fail, if e.g. the first one in a list like ['.bmp', '.jpg', ...] raises an AssertionError, because the others in the list won't be tested.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The main idea is to reducing the number of tests and also testing times. Comparing https://github.com/GenericMappingTools/pygmt/actions/runs/6677102075/job/18146781350?pr=2766 and https://github.com/GenericMappingTools/pygmt/actions/runs/6648648090/job/18065964154?pr=2766, the old way (using parametrize) has 7 tests for worldfile, and the testing time is 0.5 * 5 + 0.25 * 2 = 3 seconds, while the new way (a single test) costs 1.5 s only.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah ok, I see that we avoid having to call fig.basemap multiple times, so it will be faster. All good then.

"""
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():
"""
Expand Down