update testing.integrate_kernel type hints

Author: chrishavlin

yt/testing.py

@@ -17,6 +17,7 @@
 import numpy as np
 from more_itertools import always_iterable
 from numpy.random import RandomState
+from numpy.typing import NDArray
 from unyt.exceptions import UnitOperationError
 
 from yt._maintenance.deprecation import issue_deprecation_warning
@@ -91,7 +92,7 @@ def assert_rel_equal(a1, a2, decimals, err_msg="", verbose=True):
 
 # tested: volume integral is 1.
 def cubicspline_python(
-    x: float | np.ndarray,
+    x: float | NDArray[np.floating],
 ) -> np.ndarray:
     """
     cubic spline SPH kernel function for testing against more
@@ -118,7 +119,9 @@ def cubicspline_python(
 
 
 def integrate_kernel(
-    kernelfunc: Callable[[float], float], b: float, hsml: float
+    kernelfunc: Callable[[float | NDArray[np.floating]], float | NDArray[np.floating]],
+    b: float,
+    hsml: float,
 ) -> float:
     """
     integrates a kernel function over a line passing entirely
@@ -154,11 +157,11 @@ def integrate_kernel(
 
 
 def distancematrix(
-    pos3_i0: np.ndarray,
-    pos3_i1: np.ndarray,
+    pos3_i0: NDArray[np.floating],
+    pos3_i1: NDArray[np.floating],
     periodic: tuple[bool, bool, bool] = (True,) * 3,
-    periods: np.ndarray = _zeroperiods,
-) -> np.ndarray:
+    periods: NDArray[np.floating] = _zeroperiods,
+) -> NDArray[np.floating]:
     """
     Calculates the distances between two arrays of points.