Skip to content

Commit 4b1c9ac

Browse files
authored
Refactor _load_remote_dataset to avoid typing hints issues (#3558)
1 parent ce89ccb commit 4b1c9ac

File tree

1 file changed

+35
-43
lines changed

1 file changed

+35
-43
lines changed

pygmt/datasets/load_remote_dataset.py

Lines changed: 35 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -2,38 +2,33 @@
22
Internal function to load GMT remote datasets.
33
"""
44

5-
from __future__ import annotations
6-
7-
from typing import TYPE_CHECKING, ClassVar, Literal, NamedTuple
5+
from collections.abc import Sequence
6+
from typing import Any, Literal, NamedTuple
87

8+
import xarray as xr
99
from pygmt.clib import Session
1010
from pygmt.exceptions import GMTInvalidInput
1111
from pygmt.helpers import build_arg_list, kwargs_to_strings
1212
from pygmt.src import which
1313

14-
if TYPE_CHECKING:
15-
from collections.abc import Sequence
16-
17-
import xarray as xr
18-
1914

2015
class Resolution(NamedTuple):
2116
"""
2217
Resolution code, the available grid registrations and whether it is tiled.
2318
2419
Attributes
2520
----------
26-
code : str
21+
code
2722
The resolution code. E.g., "01d", "30m", "01s".
28-
registrations : list
23+
registrations
2924
A list of the accepted registrations for a given resolution. Can be either
3025
"pixel" or "gridline".
31-
tiled : bool
26+
tiled
3227
States if the grid is tiled, which requires an argument for ``region``.
3328
"""
3429

3530
code: str
36-
registrations: ClassVar[list] = ["gridline", "pixel"]
31+
registrations: Sequence[str] = ["gridline", "pixel"]
3732
tiled: bool = False
3833

3934

@@ -43,20 +38,20 @@ class GMTRemoteDataset(NamedTuple):
4338
4439
Attributes
4540
----------
46-
description : str
41+
description
4742
The name assigned as an attribute to the DataArray.
48-
units : str, None
43+
units
4944
The units of the values in the DataArray.
50-
resolutions : dict
45+
resolutions
5146
Dictionary of available resolution as keys and Resolution objects as values.
52-
extra_attributes : dict
47+
extra_attributes
5348
A dictionary of extra or unique attributes of the dataset.
5449
"""
5550

5651
description: str
5752
units: str | None
5853
resolutions: dict[str, Resolution]
59-
extra_attributes: dict
54+
extra_attributes: dict[str, Any]
6055

6156

6257
datasets = {
@@ -389,9 +384,8 @@ def _load_remote_dataset(
389384
The grid resolution. The suffix ``d``, ``m``, and ``s`` stand for arc-degrees,
390385
arc-minutes, and arc-seconds, respectively.
391386
region
392-
The subregion of the grid to load, in the form of a list
393-
[*xmin*, *xmax*, *ymin*, *ymax*] or a string *xmin/xmax/ymin/ymax*.
394-
Required for tiled grids.
387+
The subregion of the grid to load, in the form of a sequence [*xmin*, *xmax*,
388+
*ymin*, *ymax*] or an ISO country code. Required for tiled grids.
395389
registration
396390
Grid registration type. Either ``"pixel"`` for pixel registration or
397391
``"gridline"`` for gridline registration. Default is ``None``, where
@@ -417,41 +411,39 @@ def _load_remote_dataset(
417411

418412
# Check resolution
419413
if resolution not in dataset.resolutions:
420-
raise GMTInvalidInput(
414+
msg = (
421415
f"Invalid resolution '{resolution}' for {dataset.description} dataset. "
422416
f"Available resolutions are: {', '.join(dataset.resolutions)}."
423417
)
418+
raise GMTInvalidInput(msg)
424419
resinfo = dataset.resolutions[resolution]
425420

426421
# Check registration
427-
if registration is None:
428-
# Use gridline registration unless only pixel registration is available
429-
registration = "gridline" if "gridline" in resinfo.registrations else "pixel"
430-
elif registration in {"pixel", "gridline"}:
431-
if registration not in resinfo.registrations:
432-
raise GMTInvalidInput(
433-
f"{registration} registration is not available for the "
434-
f"{resolution} {dataset.description} dataset. Only "
435-
f"{resinfo.registrations[0]} registration is available."
422+
match registration:
423+
case None:
424+
# Use gridline registration unless only pixel registration is available
425+
reg = "g" if "gridline" in resinfo.registrations else "p"
426+
case x if x not in resinfo.registrations:
427+
msg = (
428+
f"Invalid grid registration '{registration}' for the {resolution} "
429+
f"{dataset.description} dataset. Should be either 'pixel', 'gridline' "
430+
"or None. Default is None, where a gridline-registered grid is "
431+
"returned unless only the pixel-registered grid is available."
436432
)
437-
else:
438-
raise GMTInvalidInput(
439-
f"Invalid grid registration: '{registration}', should be either 'pixel', "
440-
"'gridline' or None. Default is None, where a gridline-registered grid is "
441-
"returned unless only the pixel-registered grid is available."
442-
)
433+
raise GMTInvalidInput(msg)
434+
case _:
435+
reg = registration[0]
443436

444-
fname = f"@{prefix}_{resolution}_{registration[0]}" # type: ignore[index]
445437
if resinfo.tiled and region is None:
446-
raise GMTInvalidInput(
447-
f"'region' is required for {dataset.description} resolution '{resolution}'."
438+
msg = (
439+
f"The 'region' parameter is required for {dataset.description} "
440+
f"resolution '{resolution}'."
448441
)
442+
raise GMTInvalidInput(msg)
449443

444+
fname = f"@{prefix}_{resolution}_{reg}"
450445
kind = "image" if name in {"earth_day", "earth_night"} else "grid"
451-
kwdict = {
452-
"R": region, # region can be None
453-
"T": "i" if kind == "image" else "g",
454-
}
446+
kwdict = {"R": region, "T": {"grid": "g", "image": "i"}[kind]}
455447
with Session() as lib:
456448
with lib.virtualfile_out(kind=kind) as voutgrd:
457449
lib.call_module(

0 commit comments

Comments
 (0)