Add a type stub for the compiled IM._core extension

IM/_core.pyi

@@ -0,0 +1,29 @@
+"""Type stub for the compiled Rust extension (`src-rust/lib.rs`).
+
+Keep these signatures in step with the `#[pyfunction]`s in `src-rust/lib.rs`.
+Every array argument must be a C-contiguous array of the stated dtype.
+"""
+
+import numpy as np
+
+type _Array1F64 = np.ndarray[tuple[int], np.dtype[np.float64]]
+type _Array2F64 = np.ndarray[tuple[int, int], np.dtype[np.float64]]
+type _Array2F32 = np.ndarray[tuple[int, int], np.dtype[np.float32]]
+type _Array3F64 = np.ndarray[tuple[int, int, int], np.dtype[np.float64]]
+
+def _arias_intensity(waveforms_py: np.ndarray, dt: float) -> _Array1F64: ...
+def _cav(waveforms_py: np.ndarray, dt: float) -> _Array1F64: ...
+def _konno_ohmachi_matrix_rows(
+    n_bins: int, bandwidth: float, start: int, stop: int
+) -> _Array2F32: ...
+def _konno_ohmachi_smooth(spectra_py: np.ndarray, bandwidth: float) -> _Array2F64: ...
+def _psa(
+    comp_0_py: np.ndarray,
+    comp_90_py: np.ndarray,
+    comp_ver_py: np.ndarray,
+    coefficients_py: np.ndarray,
+) -> _Array3F64: ...
+def _rotd(comp_0_py: np.ndarray, comp_90_py: np.ndarray) -> _Array2F64: ...
+def _significant_duration(
+    waveforms_py: np.ndarray, dt: float, low: float, high: float
+) -> _Array1F64: ...

IM/ims.py

@@ -11,7 +11,7 @@
 import xarray as xr
 
 from IM import (
-    _core,  # ty: ignore[unresolved-import]
+    _core,
     konno_ohmachi,
 )
 
@@ -744,7 +744,10 @@ def _psa_kernel(
             # Keep the original samples exactly rather than their
             # interpolated round trip.
             chunk = original if factor == 1 else upsampled[..., :: max_factor // factor]
-            psa[stations, factors == factor] = _core._psa(*chunk, coefficients)
+            comp_0, comp_90, comp_ver = chunk
+            psa[stations, factors == factor] = _core._psa(
+                comp_0, comp_90, comp_ver, coefficients
+            )
     return psa.reshape(lead + (len(periods), len(ROTD_COMPONENTS)))
 
 

IM/konno_ohmachi.py

@@ -12,9 +12,7 @@
 import numpy as np
 import numpy.typing as npt
 
-from IM import (
-    _core,  # ty: ignore[unresolved-import]
-)
+from IM import _core
 
 DEFAULT_BANDWIDTH = 188.5
 """Bandwidth of the Konno-Ohmachi window. Lower values smooth more strongly.

tests/test_core_stub.py

@@ -0,0 +1,45 @@
+"""Check that the `IM._core` type stub matches the compiled extension."""
+
+import ast
+import inspect
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from IM import _core
+
+STUB_PATH = Path(_core.__file__).parent / "_core.pyi"
+
+
+def _stub_functions() -> dict[str, list[str]]:
+    tree = ast.parse(STUB_PATH.read_text())
+    return {
+        node.name: [arg.arg for arg in node.args.posonlyargs + node.args.args]
+        for node in tree.body
+        if isinstance(node, ast.FunctionDef)
+    }
+
+
+def _core_functions() -> dict[str, Callable[..., Any]]:
+    return {
+        name: value
+        for name, value in vars(_core).items()
+        if callable(value) and not name.startswith("__")
+    }
+
+
+def test_stub_exists() -> None:
+    assert STUB_PATH.is_file()
+
+
+def test_stub_declares_every_core_function() -> None:
+    assert set(_stub_functions()) == set(_core_functions())
+
+
+@pytest.mark.parametrize("name", sorted(_core_functions()))
+def test_stub_signature_matches_core(name: str) -> None:
+    function = _core_functions()[name]
+    runtime_parameters = list(inspect.signature(function).parameters)
+    assert _stub_functions().get(name) == runtime_parameters