diff --git a/clouddrift/adapters/gdp/gdpsource.py b/clouddrift/adapters/gdp/gdpsource.py index 989deb61..8c84db4c 100644 --- a/clouddrift/adapters/gdp/gdpsource.py +++ b/clouddrift/adapters/gdp/gdpsource.py @@ -484,7 +484,7 @@ def _combine_chunked_drifter_datasets(datasets: list[xr.Dataset]) -> xr.Dataset: ) sort_coord = traj_dataset.coords["obs_index"] - vals: np_typing.NDArray[np.int_] = sort_coord.data + vals: np_typing.NDArray[np.int64] = sort_coord.data sort_coord_dim = sort_coord.dims[-1] sort_key = vals.argsort() diff --git a/clouddrift/adapters/hurdat2.py b/clouddrift/adapters/hurdat2.py index da1166d3..a9f3cab3 100644 --- a/clouddrift/adapters/hurdat2.py +++ b/clouddrift/adapters/hurdat2.py @@ -390,7 +390,9 @@ def to_raggedarray( preprocess_func=lambda idx: track_data[idx].to_xarray_dataset(), attrs_global=TrackData.global_attrs, attrs_variables={ - field.name: field.metadata + field.name: dict( + field.metadata + ) # type cast needed for static type analysis for field in fields(HeaderLine) + fields(DataLine) }, ) diff --git a/clouddrift/kinematics.py b/clouddrift/kinematics.py index e105e8a4..b85ab880 100644 --- a/clouddrift/kinematics.py +++ b/clouddrift/kinematics.py @@ -2,13 +2,11 @@ Functions for kinematic computations. """ -import typing - import numpy as np import numpy.typing as np_typing -import pandas as pd import xarray as xr +import clouddrift.typing as cd_typing from clouddrift.sphere import ( EARTH_RADIUS_METERS, bearing, @@ -24,17 +22,8 @@ def kinetic_energy( - u: float - | list[typing.Any] - | np_typing.NDArray[typing.Any] - | xr.DataArray - | pd.Series, - v: float - | list - | np_typing.NDArray[typing.Any] - | xr.DataArray - | pd.Series - | None = None, + u: float | cd_typing.ArrayTypes, + v: float | cd_typing.ArrayTypes, ) -> float | np_typing.NDArray[np.float64] | xr.DataArray: """Compute kinetic energy from zonal and meridional velocities. diff --git a/clouddrift/ragged.py b/clouddrift/ragged.py index 38cc76aa..489e3421 100644 --- a/clouddrift/ragged.py +++ b/clouddrift/ragged.py @@ -460,8 +460,8 @@ def rowsize_to_index(rowsize: cd_typing.ArrayTypes) -> np.ndarray: def segment( x: cd_typing.ArrayTypes, tolerance: float | np.timedelta64 | timedelta | pd.Timedelta, - rowsize: np_typing.NDArray[np.int_] | None = None, -) -> np_typing.NDArray[np.int_]: + rowsize: np_typing.NDArray[np.int64] | None = None, +) -> np_typing.NDArray[np.int64]: """Divide an array into segments based on a tolerance value. Parameters @@ -797,7 +797,7 @@ def subset( def unpack( ragged_array: cd_typing.ArrayTypes, rowsize: cd_typing.ArrayTypes, - rows: int | np.int_ | Iterable[int] | None = None, + rows: int | np.int64 | Iterable[int] | None = None, axis: int = 0, ) -> list[np_typing.NDArray[typing.Any]]: """Unpack a ragged array into a list of regular arrays. diff --git a/clouddrift/raggedarray.py b/clouddrift/raggedarray.py index bcb4cbcf..ae1f0696 100644 --- a/clouddrift/raggedarray.py +++ b/clouddrift/raggedarray.py @@ -277,7 +277,7 @@ def from_xarray( @staticmethod def number_of_observations( rowsize_func: Callable[[int], int], indices: cd_typing.ArrayTypes, **kwargs - ) -> NDArray[np.int_]: + ) -> NDArray[np.int64]: """Iterate through the files and evaluate the number of observations. Parameters diff --git a/tests/datasets_tests.py b/tests/datasets_tests.py index 375f2281..1ab0deae 100644 --- a/tests/datasets_tests.py +++ b/tests/datasets_tests.py @@ -42,6 +42,7 @@ def test_glad_subset_and_apply_ragged_work(self): mean_lon = apply_ragged(self._mean, [ds_sub.longitude], ds_sub.rowsize) self.assertTrue(len(mean_lon) == 2) + # For static typing purposes def _mean(self, x: np_typing.NDArray[typing.Any]) -> np_typing.NDArray[typing.Any]: return np.mean(x)