TYP: Add typing hints to the return values for modules that output a grid (#3385)

seisman · web-flow · commit 819bac13a97c · 2024-08-08T22:26:26.000+08:00
diff --git a/pygmt/src/binstats.py b/pygmt/src/binstats.py
@@ -2,6 +2,7 @@
 binstats - Bin spatial data and determine statistics per bin
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
@@ -23,7 +24,7 @@
     r="registration",
 )
 @kwargs_to_strings(I="sequence", R="sequence", i="sequence_comma")
-def binstats(data, outgrid: str | None = None, **kwargs):
+def binstats(data, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Bin spatial data and determine statistics per bin.
 
@@ -95,7 +96,7 @@ def binstats(data, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/dimfilter.py b/pygmt/src/dimfilter.py
@@ -2,6 +2,7 @@
 dimfilter - Directional filtering of grids in the space domain.
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
@@ -19,7 +20,7 @@
     V="verbose",
 )
 @kwargs_to_strings(I="sequence", R="sequence")
-def dimfilter(grid, outgrid: str | None = None, **kwargs):
+def dimfilter(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Filter a grid by dividing the filter circle.
 
@@ -109,7 +110,7 @@ def dimfilter(grid, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/grdclip.py b/pygmt/src/grdclip.py
@@ -2,6 +2,7 @@
 grdclip - Change the range and extremes of grid values.
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
@@ -24,7 +25,7 @@
     Si="sequence",
     Sr="sequence",
 )
-def grdclip(grid, outgrid: str | None = None, **kwargs):
+def grdclip(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Set values in a grid that meet certain criteria to a new value.
 
@@ -62,7 +63,7 @@ def grdclip(grid, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/grdcut.py b/pygmt/src/grdcut.py
@@ -2,6 +2,7 @@
 grdcut - Extract subregion from a grid.
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -27,7 +28,7 @@
     f="coltypes",
 )
 @kwargs_to_strings(R="sequence")
-def grdcut(grid, **kwargs):
+def grdcut(grid, **kwargs) -> xr.DataArray | None:
     r"""
     Extract subregion from a grid.
 
@@ -80,7 +81,7 @@ def grdcut(grid, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/grdfill.py b/pygmt/src/grdfill.py
@@ -2,6 +2,7 @@
 grdfill - Fill blank areas from a grid.
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
@@ -17,7 +18,7 @@
     V="verbose",
 )
 @kwargs_to_strings(R="sequence")
-def grdfill(grid, outgrid: str | None = None, **kwargs):
+def grdfill(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Fill blank areas from a grid file.
 
@@ -51,7 +52,7 @@ def grdfill(grid, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/grdfilter.py b/pygmt/src/grdfilter.py
@@ -2,6 +2,7 @@
 grdfilter - Filter a grid in the space (or time) domain.
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
@@ -20,7 +21,7 @@
     x="cores",
 )
 @kwargs_to_strings(I="sequence", R="sequence")
-def grdfilter(grid, outgrid: str | None = None, **kwargs):
+def grdfilter(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Filter a grid in the space (or time) domain.
 
@@ -97,7 +98,7 @@ def grdfilter(grid, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/grdgradient.py b/pygmt/src/grdgradient.py
@@ -2,6 +2,7 @@
 grdgradient - Compute directional gradients from a grid.
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -29,7 +30,7 @@
     n="interpolation",
 )
 @kwargs_to_strings(A="sequence", E="sequence", R="sequence")
-def grdgradient(grid, outgrid: str | None = None, **kwargs):
+def grdgradient(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Compute the directional derivative of the vector gradient of the data.
 
@@ -138,7 +139,7 @@ def grdgradient(grid, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/grdhisteq.py b/pygmt/src/grdhisteq.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 import pandas as pd
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -61,7 +62,9 @@ class grdhisteq:  # noqa: N801
         h="header",
     )
     @kwargs_to_strings(R="sequence")
-    def equalize_grid(grid, outgrid: str | None = None, **kwargs):
+    def equalize_grid(
+        grid, outgrid: str | None = None, **kwargs
+    ) -> xr.DataArray | None:
         r"""
         Perform histogram equalization for a grid.
 
@@ -93,7 +96,7 @@ def equalize_grid(grid, outgrid: str | None = None, **kwargs):
 
         Returns
         -------
-        ret: xarray.DataArray or None
+        ret
             Return type depends on the ``outgrid`` parameter:
 
             - xarray.DataArray if ``outgrid`` is None
diff --git a/pygmt/src/grdlandmask.py b/pygmt/src/grdlandmask.py
@@ -2,6 +2,7 @@
 grdlandmask - Create a "wet-dry" mask grid from shoreline data base
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
@@ -22,7 +23,7 @@
     x="cores",
 )
 @kwargs_to_strings(I="sequence", R="sequence", N="sequence", E="sequence")
-def grdlandmask(outgrid: str | None = None, **kwargs):
+def grdlandmask(outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Create a grid file with set values for land and water.
 
@@ -80,7 +81,7 @@ def grdlandmask(outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/grdproject.py b/pygmt/src/grdproject.py
@@ -2,6 +2,7 @@
 grdproject - Forward and inverse map transformation of grids.
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
@@ -24,7 +25,7 @@
     r="registration",
 )
 @kwargs_to_strings(C="sequence", D="sequence", R="sequence")
-def grdproject(grid, outgrid: str | None = None, **kwargs):
+def grdproject(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Change projection of gridded data between geographical and rectangular.
 
@@ -84,7 +85,7 @@ def grdproject(grid, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/grdsample.py b/pygmt/src/grdsample.py
@@ -2,6 +2,7 @@
 grdsample - Resample a grid onto a new lattice
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
@@ -20,7 +21,7 @@
     x="cores",
 )
 @kwargs_to_strings(I="sequence", R="sequence")
-def grdsample(grid, outgrid: str | None = None, **kwargs):
+def grdsample(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Change the registration, spacing, or nodes in a grid file.
 
@@ -60,7 +61,7 @@ def grdsample(grid, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/nearneighbor.py b/pygmt/src/nearneighbor.py
@@ -2,6 +2,7 @@
 nearneighbor - Grid table data using a "Nearest neighbor" algorithm.
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
@@ -30,7 +31,7 @@
 @kwargs_to_strings(I="sequence", R="sequence", i="sequence_comma")
 def nearneighbor(
     data=None, x=None, y=None, z=None, outgrid: str | None = None, **kwargs
-):
+) -> xr.DataArray | None:
     r"""
     Grid table data using a "Nearest neighbor" algorithm.
 
@@ -117,7 +118,7 @@ def nearneighbor(
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray`: if ``outgrid`` is not set
diff --git a/pygmt/src/sph2grd.py b/pygmt/src/sph2grd.py
@@ -2,6 +2,7 @@
 sph2grd - Compute grid from spherical harmonic coefficients
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
@@ -20,7 +21,7 @@
     x="cores",
 )
 @kwargs_to_strings(I="sequence", R="sequence", i="sequence_comma")
-def sph2grd(data, outgrid: str | None = None, **kwargs):
+def sph2grd(data, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Create spherical grid files in tension of data.
 
@@ -50,7 +51,7 @@ def sph2grd(data, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/sphdistance.py b/pygmt/src/sphdistance.py
@@ -3,6 +3,7 @@
 or natural nearest-neighbor grid on a sphere
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
@@ -23,7 +24,9 @@
     V="verbose",
 )
 @kwargs_to_strings(I="sequence", R="sequence")
-def sphdistance(data=None, x=None, y=None, outgrid: str | None = None, **kwargs):
+def sphdistance(
+    data=None, x=None, y=None, outgrid: str | None = None, **kwargs
+) -> xr.DataArray | None:
     r"""
     Create Voronoi distance, node, or natural nearest-neighbor grid on a sphere.
 
@@ -86,7 +89,7 @@ def sphdistance(data=None, x=None, y=None, outgrid: str | None = None, **kwargs)
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/sphinterpolate.py b/pygmt/src/sphinterpolate.py
@@ -2,6 +2,7 @@
 sphinterpolate - Spherical gridding in tension of data on a sphere
 """
 
+import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
@@ -15,7 +16,7 @@
     V="verbose",
 )
 @kwargs_to_strings(I="sequence", R="sequence")
-def sphinterpolate(data, outgrid: str | None = None, **kwargs):
+def sphinterpolate(data, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     r"""
     Create spherical grid files in tension of data.
 
@@ -42,7 +43,7 @@ def sphinterpolate(data, outgrid: str | None = None, **kwargs):
 
     Returns
     -------
-    ret: xarray.DataArray or None
+    ret
         Return type depends on whether the ``outgrid`` parameter is set:
 
         - :class:`xarray.DataArray` if ``outgrid`` is not set
diff --git a/pygmt/src/surface.py b/pygmt/src/surface.py
diff --git a/pygmt/src/triangulate.py b/pygmt/src/triangulate.py
diff --git a/pygmt/src/xyz2grd.py b/pygmt/src/xyz2grd.py
