From 4b1c9accee5df29ade8a937eb57f6242ea83f99e Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Tue, 29 Oct 2024 22:44:38 +0800 Subject: [PATCH] Refactor _load_remote_dataset to avoid typing hints issues (#3558) --- pygmt/datasets/load_remote_dataset.py | 78 ++++++++++++--------------- 1 file changed, 35 insertions(+), 43 deletions(-) diff --git a/pygmt/datasets/load_remote_dataset.py b/pygmt/datasets/load_remote_dataset.py index 428fb68ab42..168a93583b2 100644 --- a/pygmt/datasets/load_remote_dataset.py +++ b/pygmt/datasets/load_remote_dataset.py @@ -2,20 +2,15 @@ Internal function to load GMT remote datasets. """ -from __future__ import annotations - -from typing import TYPE_CHECKING, ClassVar, Literal, NamedTuple +from collections.abc import Sequence +from typing import Any, Literal, NamedTuple +import xarray as xr from pygmt.clib import Session from pygmt.exceptions import GMTInvalidInput from pygmt.helpers import build_arg_list, kwargs_to_strings from pygmt.src import which -if TYPE_CHECKING: - from collections.abc import Sequence - - import xarray as xr - class Resolution(NamedTuple): """ @@ -23,17 +18,17 @@ class Resolution(NamedTuple): Attributes ---------- - code : str + code The resolution code. E.g., "01d", "30m", "01s". - registrations : list + registrations A list of the accepted registrations for a given resolution. Can be either "pixel" or "gridline". - tiled : bool + tiled States if the grid is tiled, which requires an argument for ``region``. """ code: str - registrations: ClassVar[list] = ["gridline", "pixel"] + registrations: Sequence[str] = ["gridline", "pixel"] tiled: bool = False @@ -43,20 +38,20 @@ class GMTRemoteDataset(NamedTuple): Attributes ---------- - description : str + description The name assigned as an attribute to the DataArray. - units : str, None + units The units of the values in the DataArray. - resolutions : dict + resolutions Dictionary of available resolution as keys and Resolution objects as values. - extra_attributes : dict + extra_attributes A dictionary of extra or unique attributes of the dataset. """ description: str units: str | None resolutions: dict[str, Resolution] - extra_attributes: dict + extra_attributes: dict[str, Any] datasets = { @@ -389,9 +384,8 @@ def _load_remote_dataset( The grid resolution. The suffix ``d``, ``m``, and ``s`` stand for arc-degrees, arc-minutes, and arc-seconds, respectively. region - The subregion of the grid to load, in the form of a list - [*xmin*, *xmax*, *ymin*, *ymax*] or a string *xmin/xmax/ymin/ymax*. - Required for tiled grids. + The subregion of the grid to load, in the form of a sequence [*xmin*, *xmax*, + *ymin*, *ymax*] or an ISO country code. Required for tiled grids. registration Grid registration type. Either ``"pixel"`` for pixel registration or ``"gridline"`` for gridline registration. Default is ``None``, where @@ -417,41 +411,39 @@ def _load_remote_dataset( # Check resolution if resolution not in dataset.resolutions: - raise GMTInvalidInput( + msg = ( f"Invalid resolution '{resolution}' for {dataset.description} dataset. " f"Available resolutions are: {', '.join(dataset.resolutions)}." ) + raise GMTInvalidInput(msg) resinfo = dataset.resolutions[resolution] # Check registration - if registration is None: - # Use gridline registration unless only pixel registration is available - registration = "gridline" if "gridline" in resinfo.registrations else "pixel" - elif registration in {"pixel", "gridline"}: - if registration not in resinfo.registrations: - raise GMTInvalidInput( - f"{registration} registration is not available for the " - f"{resolution} {dataset.description} dataset. Only " - f"{resinfo.registrations[0]} registration is available." + match registration: + case None: + # Use gridline registration unless only pixel registration is available + reg = "g" if "gridline" in resinfo.registrations else "p" + case x if x not in resinfo.registrations: + msg = ( + f"Invalid grid registration '{registration}' for the {resolution} " + f"{dataset.description} dataset. Should be either 'pixel', 'gridline' " + "or None. Default is None, where a gridline-registered grid is " + "returned unless only the pixel-registered grid is available." ) - else: - raise GMTInvalidInput( - f"Invalid grid registration: '{registration}', should be either 'pixel', " - "'gridline' or None. Default is None, where a gridline-registered grid is " - "returned unless only the pixel-registered grid is available." - ) + raise GMTInvalidInput(msg) + case _: + reg = registration[0] - fname = f"@{prefix}_{resolution}_{registration[0]}" # type: ignore[index] if resinfo.tiled and region is None: - raise GMTInvalidInput( - f"'region' is required for {dataset.description} resolution '{resolution}'." + msg = ( + f"The 'region' parameter is required for {dataset.description} " + f"resolution '{resolution}'." ) + raise GMTInvalidInput(msg) + fname = f"@{prefix}_{resolution}_{reg}" kind = "image" if name in {"earth_day", "earth_night"} else "grid" - kwdict = { - "R": region, # region can be None - "T": "i" if kind == "image" else "g", - } + kwdict = {"R": region, "T": {"grid": "g", "image": "i"}[kind]} with Session() as lib: with lib.virtualfile_out(kind=kind) as voutgrd: lib.call_module(