From fdf817e4f743e444528a22501eec5854c6a968c6 Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Tue, 29 Oct 2024 13:01:11 +0800 Subject: [PATCH] Move the parser function into _common.py --- pygmt/src/_common.py | 59 +++++++++++++++++++++++++++++++++++++++- pygmt/src/coast.py | 12 ++++---- pygmt/src/grdlandmask.py | 7 +++-- pygmt/src/select.py | 5 ++-- 4 files changed, 71 insertions(+), 12 deletions(-) diff --git a/pygmt/src/_common.py b/pygmt/src/_common.py index 96ae9791b71..b3842cc5cf9 100644 --- a/pygmt/src/_common.py +++ b/pygmt/src/_common.py @@ -3,8 +3,9 @@ """ from pathlib import Path -from typing import Any +from typing import Any, Literal +from pygmt.exceptions import GMTInvalidInput from pygmt.src.which import which @@ -39,3 +40,59 @@ def _data_geometry_is_point(data: Any, kind: str) -> bool: except FileNotFoundError: pass return False + + +def _parse_coastline_resolution( + resolution: Literal["auto", "full", "high", "intermediate", "low", "crude", None], + allow_auto: bool = False, +) -> str | None: + """ + Parse the resolution parameter for coastline-related functions. + + Parameters + ---------- + resolution + The resolution of the coastline dataset to use. The available resolutions from + highest to lowest are: ``"full"``, ``"high"``, ``"intermediate"``, ``"low"``, + and ``"crude"``, which drops by 80% between levels. + allow_auto + Whether to allow the ``"auto"`` resolution. + + Returns + ------- + str or None + The parsed resolution value. + + Raises + ------ + GMTInvalidInput + If the resolution is invalid. + + Examples + -------- + >>> _parse_coastline_resolution("full") + "f" + >>> _parse_coastline_resolution("f") + "f" + >>> _parse_coastline_resolution("auto", allow_auto=True) + "a" + >>> _parse_coastline_resolution("invalid") + pygmt.exceptions.GMTInvalidInput: Invalid resolution: invalid. Valid values are ... + >>> _parse_coastline_resolution(None) + None + >>> _parse_coastline_resolution("auto") + pygmt.exceptions.GMTInvalidInput: Invalid resolution: auto. Valid values are ... + """ + if resolution is None: + return None + + valid_resolutions = {"full", "high", "intermediate", "low", "crude"} + if allow_auto: + valid_resolutions.add("auto") + if resolution not in {*valid_resolutions, *[res[0] for res in valid_resolutions]}: + msg = ( + f"Invalid resolution: {resolution}." + f"Valid values are {', '.join(valid_resolutions)}." + ) + raise GMTInvalidInput(msg) + return resolution[0] diff --git a/pygmt/src/coast.py b/pygmt/src/coast.py index fe0a034ec30..1c5cbfc147f 100644 --- a/pygmt/src/coast.py +++ b/pygmt/src/coast.py @@ -13,6 +13,7 @@ kwargs_to_strings, use_alias, ) +from pygmt.src._common import _parse_coastline_resolution __doctest_skip__ = ["coast"] @@ -41,8 +42,8 @@ def coast( self, resolution: Literal[ - "auto", "full", "high", "intermediate", "low", "crude" - ] = "auto", + "auto", "full", "high", "intermediate", "low", "crude", None + ] = None, **kwargs, ): r""" @@ -208,9 +209,8 @@ def coast( """At least one of the following parameters must be specified: lakes, land, water, rivers, borders, dcw, Q, or shorelines""" ) - - # Alias 'resolution' to "D". - kwargs["D"] = resolution[0] - + kwargs["D"] = kwargs.get( + "D", _parse_coastline_resolution(resolution, allow_auto=True) + ) with Session() as lib: lib.call_module(module="coast", args=build_arg_list(kwargs)) diff --git a/pygmt/src/grdlandmask.py b/pygmt/src/grdlandmask.py index b99b13497ad..a654a3b4e62 100644 --- a/pygmt/src/grdlandmask.py +++ b/pygmt/src/grdlandmask.py @@ -8,6 +8,7 @@ from pygmt.clib import Session from pygmt.exceptions import GMTInvalidInput from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias +from pygmt.src._common import _parse_coastline_resolution __doctest_skip__ = ["grdlandmask"] @@ -26,7 +27,7 @@ @kwargs_to_strings(I="sequence", R="sequence", N="sequence", E="sequence") def grdlandmask( outgrid: str | None = None, - resolution: Literal["full", "high", "intermediate", "low", "crude"] = "low", + resolution: Literal["full", "high", "intermediate", "low", "crude", None] = None, **kwargs, ) -> xr.DataArray | None: r""" @@ -99,8 +100,8 @@ def grdlandmask( if kwargs.get("I") is None or kwargs.get("R") is None: raise GMTInvalidInput("Both 'region' and 'spacing' must be specified.") - # Alias "resolution" to "D" - kwargs["D"] = resolution[0] + kwargs["D"] = kwargs.get("D", _parse_coastline_resolution(resolution)) + with Session() as lib: with lib.virtualfile_out(kind="grid", fname=outgrid) as voutgrd: kwargs["G"] = voutgrd diff --git a/pygmt/src/select.py b/pygmt/src/select.py index 21d2c856217..a5bef0d3556 100644 --- a/pygmt/src/select.py +++ b/pygmt/src/select.py @@ -14,6 +14,7 @@ use_alias, validate_output_table_type, ) +from pygmt.src._common import _parse_coastline_resolution __doctest_skip__ = ["select"] @@ -47,7 +48,7 @@ def select( data=None, output_type: Literal["pandas", "numpy", "file"] = "pandas", outfile: str | None = None, - resolution: Literal["full", "high", "intermediate", "low", "crude"] = "low", + resolution: Literal["full", "high", "intermediate", "low", "crude", None] = None, **kwargs, ) -> pd.DataFrame | np.ndarray | None: r""" @@ -203,7 +204,7 @@ def select( >>> out = pygmt.select(data=ship_data, region=[246, 247, 20, 21]) """ # Alias "resolution" to "D" - kwargs["D"] = resolution[0] + kwargs["D"] = kwargs.get("D", _parse_coastline_resolution(resolution)) output_type = validate_output_table_type(output_type, outfile=outfile)