Refine typing with NumPy plugin of mypy (#385)

* Make numpy plugin of mypy happy

* Finish linting!

* Add comment for backward comp

* Explicitly define union

* Exception for Number before Python 3.9

* Clarify backward behaviour

* Linting

* Try with union everywhere

* Fix dtype union

* Fix bool type

* Fix bool again

* Temporarily ignore boolean inconsistencies with Raster

* Remove RGB example in colorbar test

* Add comment

* Linting

Author: rhugonnet

.pre-commit-config.yaml

@@ -48,7 +48,9 @@ repos:
           hooks:
                 - id: mypy
                   args: [
+                        --config-file=mypy.ini,
                         --strict,
+                        --implicit-optional,
                         --ignore-missing-imports,  # Don't warn about stubs since pre-commit runs in a limited env
                         --allow-untyped-calls,  # Dynamic function/method calls are okay. Untyped function definitions are not okay.
                         --show-error-codes,
@@ -58,7 +60,7 @@ repos:
                         --disable-error-code=var-annotated,
                         --disable-error-code=no-any-return
                   ]
-                  additional_dependencies: [tokenize-rt==3.2.0]
+                  additional_dependencies: [tokenize-rt==3.2.0, numpy==1.22]
                   files: ^(geoutils|tests)
 
         # Sort imports using isort

geoutils/_typing.py

@@ -1,14 +1,46 @@
 """Typing aliases for internal use."""
 from __future__ import annotations
 
+import sys
 from typing import Any, List, Tuple, Union
 
 import numpy as np
 
-# Make an array-like type (since the array-like numpy type only exists in numpy>=1.20)
-ArrayLike = Union[np.ndarray, np.ma.masked_array, List[Any], Tuple[Any]]
+# Only for Python >= 3.9
+if sys.version_info.minor >= 9:
 
-DTypeLike = Union[str, type, np.dtype]
+    from numpy.typing import (  # this syntax works starting on Python 3.9
+        ArrayLike,
+        DTypeLike,
+        NDArray,
+    )
 
-# Mypy has issues with the builtin Number type (https://github.com/python/mypy/issues/3186)
-AnyNumber = Union[int, float, np.number]
+    # Mypy has issues with the builtin Number type (https://github.com/python/mypy/issues/3186)
+    Number = Union[int, float, np.integer[Any], np.floating[Any]]
+
+    # Simply define here if they exist
+    DTypeLike = DTypeLike
+    ArrayLike = ArrayLike
+
+    # Use NDArray wrapper to easily define numerical (float or int) N-D array types, and boolean N-D array types
+    NDArrayNum = NDArray[Union[np.floating[Any], np.integer[Any]]]
+    NDArrayBool = NDArray[np.bool_]
+    # Define numerical (float or int) masked N-D array type
+    MArrayNum = np.ma.masked_array[Any, np.dtype[Union[np.floating[Any], np.integer[Any]]]]
+    MArrayBool = np.ma.masked_array[Any, np.dtype[np.bool_]]
+
+# For backward compatibility before Python 3.9
+else:
+
+    # Mypy has issues with the builtin Number type (https://github.com/python/mypy/issues/3186)
+    Number = Union[int, float, np.integer, np.floating]  # type: ignore
+
+    # Make an array-like type (since the array-like numpy type only exists in numpy>=1.20)
+    DTypeLike = Union[str, type, np.dtype]  # type: ignore
+    ArrayLike = Union[np.ndarray, np.ma.masked_array, List[Any], Tuple[Any]]  # type: ignore
+
+    # Define generic types for NumPy array and masked-array (behaves as "Any" before 3.9 and plugin)
+    NDArrayNum = np.ndarray  # type: ignore
+    NDArrayBool = np.ndarray  # type: ignore
+    MArrayNum = np.ma.masked_array  # type: ignore
+    MArrayBool = np.ma.masked_array  # type: ignore

geoutils/projtools.py

@@ -8,7 +8,6 @@
 from math import ceil, floor
 from typing import Literal
 
-import geopandas
 import geopandas as gpd
 import numpy as np
 import pyproj
@@ -18,8 +17,10 @@
 from shapely.geometry.base import BaseGeometry
 from shapely.geometry.polygon import Polygon
 
+from geoutils._typing import NDArrayNum, Number
 
-def latlon_to_utm(lat: float, lon: float) -> str:
+
+def latlon_to_utm(lat: Number, lon: Number) -> str:
     """
     Get UTM zone for a given latitude and longitude coordinates.
 
@@ -108,37 +109,37 @@ def _get_utm_ups_crs(df: gpd.GeoDataFrame, method: Literal["centroid"] | Literal
 
 
 def bounds2poly(
-    boundsGeom: list[float] | rio.io.DatasetReader,
+    bounds_geom: list[float] | rio.io.DatasetReader,
     in_crs: CRS | None = None,
     out_crs: CRS | None = None,
 ) -> Polygon:
     """
     Converts self's bounds into a shapely Polygon. Optionally, returns it into a different CRS.
 
-    :param boundsGeom: A geometry with bounds. Can be either a list of coordinates (xmin, ymin, xmax, ymax),\
+    :param bounds_geom: A geometry with bounds. Can be either a list of coordinates (xmin, ymin, xmax, ymax),\
             a rasterio/Raster object, a geoPandas/Vector object
     :param in_crs: Input CRS
     :param out_crs: Output CRS
 
     :returns: Output polygon
     """
     # If boundsGeom is a GeoPandas or Vector object (warning, has both total_bounds and bounds attributes)
-    if hasattr(boundsGeom, "total_bounds"):
-        xmin, ymin, xmax, ymax = boundsGeom.total_bounds  # type: ignore
-        in_crs = boundsGeom.crs  # type: ignore
+    if hasattr(bounds_geom, "total_bounds"):
+        xmin, ymin, xmax, ymax = bounds_geom.total_bounds  # type: ignore
+        in_crs = bounds_geom.crs  # type: ignore
     # If boundsGeom is a rasterio or Raster object
-    elif hasattr(boundsGeom, "bounds"):
-        xmin, ymin, xmax, ymax = boundsGeom.bounds  # type: ignore
-        in_crs = boundsGeom.crs  # type: ignore
+    elif hasattr(bounds_geom, "bounds"):
+        xmin, ymin, xmax, ymax = bounds_geom.bounds  # type: ignore
+        in_crs = bounds_geom.crs  # type: ignore
     # if a list of coordinates
-    elif isinstance(boundsGeom, (list, tuple)):
-        xmin, ymin, xmax, ymax = boundsGeom
+    elif isinstance(bounds_geom, (list, tuple)):
+        xmin, ymin, xmax, ymax = bounds_geom
     else:
         raise ValueError(
             "boundsGeom must a list/tuple of coordinates or an object with attributes bounds or total_bounds."
         )
 
-    corners = ((xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax))
+    corners = np.array([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
 
     if (in_crs is not None) & (out_crs is not None):
         corners = np.transpose(reproject_points(np.transpose(corners), in_crs, out_crs))
@@ -176,7 +177,7 @@ def merge_bounds(
             "or total_bounds"
         )
 
-    output_poly = bounds2poly(boundsGeom=bounds_list[0])
+    output_poly = bounds2poly(bounds_geom=bounds_list[0])
 
     # Compute the merging
     for boundsGeom in bounds_list[1:]:
@@ -238,7 +239,9 @@ def align_bounds(
     return (left, bottom, right, top)
 
 
-def reproject_points(pts: list[list[float]] | np.ndarray, in_crs: CRS, out_crs: CRS) -> tuple[list[float], list[float]]:
+def reproject_points(
+    pts: list[list[float]] | tuple[list[float], list[float]] | NDArrayNum, in_crs: CRS, out_crs: CRS
+) -> tuple[list[float], list[float]]:
     """
     Reproject a set of point from input_crs to output_crs.
 
@@ -262,7 +265,7 @@ def reproject_points(pts: list[list[float]] | np.ndarray, in_crs: CRS, out_crs:
 
 
 def reproject_to_latlon(
-    pts: list[list[float]] | np.ndarray, in_crs: CRS, round_: int = 8
+    pts: list[list[float]] | NDArrayNum, in_crs: CRS, round_: int = 8
 ) -> tuple[list[float], list[float]]:
     """
     Reproject a set of point from in_crs to lat/lon.
@@ -279,7 +282,7 @@ def reproject_to_latlon(
 
 
 def reproject_from_latlon(
-    pts: list[list[float]] | tuple[list[float], list[float]] | np.ndarray, out_crs: CRS, round_: int = 2
+    pts: list[list[float]] | tuple[list[float], list[float]] | NDArrayNum, out_crs: CRS, round_: int = 2
 ) -> tuple[list[float], list[float]]:
     """
     Reproject a set of point from lat/lon to out_crs.

geoutils/raster/array.py

@@ -7,9 +7,10 @@
 import numpy as np
 
 import geoutils as gu
+from geoutils._typing import MArrayNum, NDArrayNum
 
 
-def get_mask(array: np.ndarray | np.ma.masked_array) -> np.ndarray:
+def get_mask(array: NDArrayNum | MArrayNum) -> NDArrayNum:
     """
     Return the mask of invalid values, whether array is a ndarray with NaNs or a np.ma.masked_array.
 
@@ -22,8 +23,8 @@ def get_mask(array: np.ndarray | np.ma.masked_array) -> np.ndarray:
 
 
 def get_array_and_mask(
-    array: np.ndarray | np.ma.masked_array, check_shape: bool = True, copy: bool = True
-) -> tuple[np.ndarray, np.ndarray]:
+    array: NDArrayNum | MArrayNum, check_shape: bool = True, copy: bool = True
+) -> tuple[NDArrayNum, NDArrayNum]:
     """
     Return array with masked values set to NaN and the associated mask.
     Works whether array is a ndarray with NaNs or a np.ma.masked_array.
@@ -65,7 +66,7 @@ def get_array_and_mask(
     return array_data, invalid_mask
 
 
-def get_valid_extent(array: np.ndarray | np.ma.masked_array) -> tuple[int, ...]:
+def get_valid_extent(array: NDArrayNum | MArrayNum) -> tuple[int, ...]:
     """
     Return (rowmin, rowmax, colmin, colmax), the first/last row/column of array with valid pixels
     """
@@ -78,7 +79,7 @@ def get_valid_extent(array: np.ndarray | np.ma.masked_array) -> tuple[int, ...]:
     return rows_nonzero[0], rows_nonzero[-1], cols_nonzero[0], cols_nonzero[-1]
 
 
-def get_xy_rotated(raster: gu.Raster, along_track_angle: float) -> tuple[np.ndarray, np.ndarray]:
+def get_xy_rotated(raster: gu.Raster, along_track_angle: float) -> tuple[NDArrayNum, NDArrayNum]:
     """
     Rotate x, y axes of image to get along- and cross-track distances.
     :param raster: Raster to get x,y positions from.

geoutils/raster/multiraster.py

@@ -10,6 +10,7 @@
 from tqdm import tqdm
 
 import geoutils as gu
+from geoutils._typing import NDArrayNum
 from geoutils.misc import resampling_method_from_str
 from geoutils.raster import Raster, RasterType, get_array_and_mask
 from geoutils.raster.raster import _default_nodata
@@ -167,7 +168,7 @@ def stack_rasters(
         )
 
     # Make a data list and add all of the reprojected rasters into it.
-    data: list[np.ndarray] = []
+    data: list[NDArrayNum] = []
 
     for raster in tqdm(rasters, disable=not progress):
         # Check that data is loaded, otherwise temporarily load it
@@ -209,7 +210,7 @@ def stack_rasters(
         nodata = reference_raster.nodata
     else:
         nodata = _default_nodata(data.dtype)
-    data[np.isnan(data)] = nodata
+    data[np.isnan(data)] = nodata  # type: ignore
 
     # Save as gu.Raster - needed as some child classes may not accept multiple bands
     r = gu.Raster.from_array(

geoutils/raster/raster.py

@@ -2,17 +2,52 @@
 
 from __future__ import annotations
 
+from typing import Literal, overload
+
 import numpy as np
 
+from geoutils._typing import MArrayNum, NDArrayNum
 from geoutils.raster.array import get_mask
 
 
+@overload
+def subsample_array(
+    array: NDArrayNum | MArrayNum,
+    subsample: float | int,
+    return_indices: Literal[False] = False,
+    *,
+    random_state: np.random.RandomState | int | None = None,
+) -> NDArrayNum:
+    ...
+
+
+@overload
+def subsample_array(
+    array: NDArrayNum | MArrayNum,
+    subsample: float | int,
+    return_indices: Literal[True],
+    *,
+    random_state: np.random.RandomState | int | None = None,
+) -> tuple[NDArrayNum, ...]:
+    ...
+
+
+@overload
+def subsample_array(
+    array: NDArrayNum | MArrayNum,
+    subsample: float | int,
+    return_indices: bool = False,
+    random_state: np.random.RandomState | int | None = None,
+) -> NDArrayNum | tuple[NDArrayNum, ...]:
+    ...
+
+
 def subsample_array(
-    array: np.ndarray | np.ma.masked_array,
+    array: NDArrayNum | MArrayNum,
     subsample: float | int,
     return_indices: bool = False,
-    random_state: np.random.RandomState | np.random.Generator | int | None = None,
-) -> np.ndarray:
+    random_state: np.random.RandomState | int | None = None,
+) -> NDArrayNum | tuple[NDArrayNum, ...]:
     """
     Randomly subsample a 1D or 2D array by a subsampling factor, taking only non NaN/masked values.
 
@@ -26,8 +61,8 @@ def subsample_array(
     """
     # Define state for random subsampling (to fix results during testing)
     if random_state is None:
-        rnd = np.random.default_rng()
-    elif isinstance(random_state, (np.random.RandomState, np.random.Generator)):
+        rnd: np.random.RandomState | np.random.Generator = np.random.default_rng()
+    elif isinstance(random_state, np.random.RandomState):
         rnd = random_state
     else:
         rnd = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(random_state)))
@@ -98,7 +133,7 @@ def _get_closest_rectangle(size: int) -> tuple[int, int]:
     raise NotImplementedError(f"Function criteria not met for rectangle of size: {size}")
 
 
-def subdivide_array(shape: tuple[int, ...], count: int) -> np.ndarray:
+def subdivide_array(shape: tuple[int, ...], count: int) -> NDArrayNum:
     """
     Create indices for subdivison of an array in a number of blocks.
 

geoutils/raster/sampling.py

@@ -13,6 +13,7 @@
 import numpy as np
 import rasterio as rio
 
+from geoutils._typing import NDArrayNum
 from geoutils.raster import Raster, RasterType
 
 lsat_sensor = {"C": "OLI/TIRS", "E": "ETM+", "T": "TM", "M": "MSS", "O": "OLI", "TI": "TIRS"}
@@ -427,7 +428,7 @@ def __parse_metadata_from_file(self, fn_meta: str | None) -> None:
 
         return None
 
-    def copy(self, new_array: np.ndarray | None = None) -> SatelliteImage:
+    def copy(self, new_array: NDArrayNum | None = None) -> SatelliteImage:
         new_satimg = super().copy(new_array=new_array)  # type: ignore
         # all objects here are immutable so no need for a copy method (string and datetime)
         # satimg_attrs = ['satellite', 'sensor', 'product', 'version', 'tile_name', 'datetime'] #taken outside of class

geoutils/raster/satimg.py

@@ -38,6 +38,7 @@
 from shapely.geometry.polygon import Polygon
 
 import geoutils as gu
+from geoutils._typing import NDArrayBool, NDArrayNum
 from geoutils.misc import copy_doc
 from geoutils.projtools import (
     _get_bounds_projected,
@@ -1092,7 +1093,7 @@ def create_mask(
         xres: float | None = None,
         yres: float | None = None,
         bounds: tuple[float, float, float, float] | None = None,
-        buffer: int | float | np.number = 0,
+        buffer: int | float | np.integer[Any] | np.floating[Any] = 0,
         *,
         as_array: Literal[False] = False,
     ) -> gu.Mask:
@@ -1106,10 +1107,10 @@ def create_mask(
         xres: float | None = None,
         yres: float | None = None,
         bounds: tuple[float, float, float, float] | None = None,
-        buffer: int | float | np.number = 0,
+        buffer: int | float | np.integer[Any] | np.floating[Any] = 0,
         *,
         as_array: Literal[True],
-    ) -> np.ndarray:
+    ) -> NDArrayNum:
         ...
 
     def create_mask(
@@ -1119,9 +1120,9 @@ def create_mask(
         xres: float | None = None,
         yres: float | None = None,
         bounds: tuple[float, float, float, float] | None = None,
-        buffer: int | float | np.number = 0,
+        buffer: int | float | np.integer[Any] | np.floating[Any] = 0,
         as_array: bool = False,
-    ) -> gu.Mask | np.ndarray:
+    ) -> gu.Mask | NDArrayBool:
         """
         Create a mask from the vector features.
 

geoutils/vector.py

@@ -0,0 +1,2 @@
+[mypy]
+plugins = numpy.typing.mypy_plugin

mypy.ini

@@ -130,7 +130,7 @@ def test_latlon_reproject(self, example: str) -> None:
         randx = np.random.randint(low=img.bounds.left, high=img.bounds.right, size=(nsample,))
         randy = np.random.randint(low=img.bounds.bottom, high=img.bounds.top, size=(nsample,))
 
-        lat, lon = pt.reproject_to_latlon([randx, randy], img.crs)
+        lat, lon = pt.reproject_to_latlon([list(randx), list(randy)], img.crs)
         x, y = pt.reproject_from_latlon([lat, lon], img.crs)
 
         assert np.all(x == randx)

tests/test_projtools.py

@@ -20,6 +20,7 @@
 import geoutils as gu
 import geoutils.projtools as pt
 from geoutils import examples
+from geoutils._typing import MArrayNum, NDArrayNum
 from geoutils.misc import resampling_method_from_str
 from geoutils.projtools import reproject_to_latlon
 from geoutils.raster.raster import _default_nodata, _default_rio_attrs
@@ -29,7 +30,7 @@
 
 def run_gdal_proximity(
     input_raster: gu.Raster, target_values: list[float] | None, distunits: str = "GEO"
-) -> np.ndarray:
+) -> NDArrayNum:
     """Run GDAL's ComputeProximity and return the read numpy array."""
     # Rasterio strongly recommends against importing gdal along rio, so this is done here instead.
     from osgeo import gdal, gdalconst
@@ -412,8 +413,8 @@ def test_data_setter(self, dtype: str, nodata_init: str | None) -> None:
 
         # Create random values between the lower and upper limit of the data type, max absolute 99999 for floats
         if "int" in dtype:
-            val_min = np.iinfo(int_type=dtype).min
-            val_max = np.iinfo(int_type=dtype).max
+            val_min: int = np.iinfo(int_type=dtype).min  # type: ignore
+            val_max: int = np.iinfo(int_type=dtype).max  # type: ignore
             randint_dtype = dtype
         else:
             val_min = -99999
@@ -422,7 +423,9 @@ def test_data_setter(self, dtype: str, nodata_init: str | None) -> None:
 
         # Fix the random seed
         np.random.seed(42)
-        arr = np.random.randint(low=val_min, high=val_max, size=(width, height), dtype=randint_dtype).astype(dtype)
+        arr = np.random.randint(
+            low=val_min, high=val_max, size=(width, height), dtype=randint_dtype  # type: ignore
+        ).astype(dtype)
         mask = np.random.randint(0, 2, size=(width, height), dtype=bool)
 
         # Check that we are actually masking stuff
@@ -1675,7 +1678,7 @@ def test_xy2ij_and_interp(self) -> None:
         assert img[int(i), int(j)] == val
 
         # Finally, check that interp convert to latlon
-        lat, lon = gu.projtools.reproject_to_latlon((x, y), in_crs=r.crs)
+        lat, lon = gu.projtools.reproject_to_latlon([[x], [y]], in_crs=r.crs)
         val_latlon = r.interp_points([(lat, lon)], order=1, input_latlon=True)[0]
         assert val == pytest.approx(val_latlon, abs=0.0001)
 
@@ -1718,7 +1721,7 @@ def test_value_at_coords(self) -> None:
         # -- Tests 2: check arguments work as intended --
 
         # 1/ Lat-lon argument check by getting the coordinates of our last test point
-        lat, lon = reproject_to_latlon(pts=[xtest0, ytest0], in_crs=r.crs)
+        lat, lon = reproject_to_latlon(pts=[[xtest0], [ytest0]], in_crs=r.crs)
         z_val_2 = r.value_at_coords(lon, lat, latlon=True)
         assert z_val == z_val_2
 
@@ -2039,13 +2042,13 @@ def test_default_nodata(self) -> None:
 
         # Check it works with most frequent np.dtypes too
         assert _default_nodata(np.dtype("uint8")) == np.iinfo("uint8").max
-        for dtype in [np.dtype("int32"), np.dtype("float32"), np.dtype("float64")]:
-            assert _default_nodata(dtype) == -99999
+        for dtype_obj in [np.dtype("int32"), np.dtype("float32"), np.dtype("float64")]:
+            assert _default_nodata(dtype_obj) == -99999  # type: ignore
 
         # Check it works with most frequent types too
         assert _default_nodata(np.uint8) == np.iinfo("uint8").max
-        for dtype in [np.int32, np.float32, np.float64]:
-            assert _default_nodata(dtype) == -99999
+        for dtype_obj in [np.int32, np.float32, np.float64]:
+            assert _default_nodata(dtype_obj) == -99999
 
         # Check that an error is raised for other types
         expected_message = "No default nodata value set for dtype."
@@ -2105,9 +2108,8 @@ def test_astype(self) -> None:
             assert np.dtype(rout.dtypes[0]) == dtype
             assert rout.data.dtype == dtype
 
-    @pytest.mark.parametrize(
-        "example", [landsat_b4_path, landsat_b4_crop_path, landsat_rgb_path, aster_dem_path]
-    )  # type: ignore
+    # The multi-band example will not have a colorbar, so not used in tests
+    @pytest.mark.parametrize("example", [landsat_b4_path, landsat_b4_crop_path, aster_dem_path])  # type: ignore
     @pytest.mark.parametrize("figsize", np.arange(2, 20, 2))  # type: ignore
     def test_show_cbar(self, example, figsize) -> None:
         """
@@ -3216,7 +3218,7 @@ def test_reflectivity(self, ops: list[str]) -> None:
     @classmethod
     def from_array(
         cls: type[TestArithmetic],
-        data: np.ndarray | np.ma.masked_array,
+        data: NDArrayNum | MArrayNum,
         rst_ref: gu.RasterType,
         nodata: int | float | list[int] | list[float] | None = None,
     ) -> gu.Raster:
@@ -3595,7 +3597,7 @@ def test_array_ufunc_1nin_1nout(self, ufunc_str: str, nodata_init: None | str, d
             warnings.filterwarnings("ignore", category=RuntimeWarning)
 
             # Check if our input dtype is possible on this ufunc, if yes check that outputs are identical
-            if com_dtype in (str(np.dtype(t[0])) for t in ufunc.types):
+            if com_dtype in [str(np.dtype(t[0])) for t in ufunc.types]:  # noqa
                 # For a single output
                 if ufunc.nout == 1:
                     assert np.ma.allequal(ufunc(rst.data), ufunc(rst).data)
@@ -3680,7 +3682,7 @@ def test_array_ufunc_2nin_1nout(
             warnings.filterwarnings("ignore", category=UserWarning)
 
             # Check if both our input dtypes are possible on this ufunc, if yes check that outputs are identical
-            if com_dtype_tuple in ((str(np.dtype(t[0])), str(np.dtype(t[1]))) for t in ufunc.types):
+            if com_dtype_tuple in [(np.dtype(t[0]), np.dtype(t[1])) for t in ufunc.types]:  # noqa
                 # For a single output
                 if ufunc.nout == 1:
                     # There exists a single exception due to negative integers as exponent of integers in "power"

tests/test_raster.py

@@ -6,6 +6,7 @@
 import pytest
 
 import geoutils as gu
+from geoutils._typing import NDArrayNum
 
 
 class TestSubsampling:
@@ -33,7 +34,7 @@ class TestSubsampling:
     assert np.count_nonzero(array3D.mask) > 0
 
     @pytest.mark.parametrize("array", [array1D, array2D, array3D])  # type: ignore
-    def test_subsample(self, array: np.ndarray) -> None:
+    def test_subsample(self, array: NDArrayNum) -> None:
         """
         Test gu.raster.subsample_array.
         """