Skip to content

Commit

Permalink
Refactor _load_remote_dataset to avoid typing hints issues (#3558)
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman authored Oct 29, 2024
1 parent ce89ccb commit 4b1c9ac
Showing 1 changed file with 35 additions and 43 deletions.
78 changes: 35 additions & 43 deletions pygmt/datasets/load_remote_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,38 +2,33 @@
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):
"""
Resolution code, the available grid registrations and whether it is tiled.
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


Expand All @@ -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 = {
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand Down

0 comments on commit 4b1c9ac

Please sign in to comment.