simplify, cleanup and docs

Cloud-Drift · Oct 9, 2024 · 1a61e41 · 1a61e41
1 parent 69a9fe4
commit 1a61e41
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 20 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)