From 1a28ff3febf40367caed462836edb2b4b426967e Mon Sep 17 00:00:00 2001
From: "ucgmsim[bot]" <335857963+ucgmsim[bot]@users.noreply.github.com>
Date: Tue, 6 Oct 2026 02:26:35 +1300
Subject: [PATCH 1/2] Add a type stub for the compiled IM._core extension
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Closes #242

## What was wrong

The Rust extension `IM._core` had no `.pyi` stub. Both `IM/ims.py` and `IM/konno_ohmachi.py` imported it with `# ty: ignore[unresolved-import]`, so `ty` treated `_core` as unknown and never checked any call into the core.

## What changed

- **`IM/_core.pyi`** (new): signatures for all seven `#[pyfunction]`s in `src-rust/lib.rs` (`_arias_intensity`, `_cav`, `_konno_ohmachi_matrix_rows`, `_konno_ohmachi_smooth`, `_psa`, `_rotd`, `_significant_duration`). Parameter names match the runtime ones, and return types give the dtype and number of dimensions (`float32` for the Konno-Ohmachi matrix rows, `float64` everywhere else).
- **`IM/ims.py`, `IM/konno_ohmachi.py`**: removed the `ty: ignore[unresolved-import]` comments.
- **`IM/ims.py` (`_psa` call)**: once the stub was in place, `ty` flagged `_core._psa(*chunk, coefficients)`, because it can't tell that `chunk` unpacks into exactly three components. I changed it to unpack explicitly first (`comp_0, comp_90, comp_ver = chunk`). Nothing it computes changes.

## Tests

`tests/test_core_stub.py` checks that the stub file exists, that it declares exactly the functions the compiled `_core` exports, and that each function's parameter names in the stub match `inspect.signature` on the real extension. Before the change all 9 tests failed because there was no stub. They now pass, and will fail if someone adds or renames a `#[pyfunction]` without updating the stub. `.ai/ci-checks` passes in full (ruff, ty, deptry, npdlint, pytest, cargo).

## For reviewers

- **Input types are loose on purpose.** Array arguments are typed as plain `np.ndarray`, not shape- or dtype-specific arrays. At runtime PyO3 enforces C-contiguous 2-D `float64`, and the stub's docstring says so. Tighter input types would mean adding casts at every call site, which is a bigger change than this issue asks for.
- **Packaging:** the stub sits next to the compiled extension inside the `IM` package. With setuptools-scm and the default `include-package-data`, it should go into wheels once committed. No `py.typed` marker was added, since the issue didn't ask for one.

_Cost $0.45, 23 turns._

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 IM/_core.pyi            | 29 ++++++++++++++++++++++++++
 IM/ims.py               |  7 +++++--
 IM/konno_ohmachi.py     |  4 +---
 tests/test_core_stub.py | 45 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 IM/_core.pyi
 create mode 100644 tests/test_core_stub.py

diff --git a/IM/_core.pyi b/IM/_core.pyi
new file mode 100644
index 00000000..e2734074
--- /dev/null
+++ b/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
+
+_Array1F64 = np.ndarray[tuple[int], np.dtype[np.float64]]
+_Array2F64 = np.ndarray[tuple[int, int], np.dtype[np.float64]]
+_Array2F32 = np.ndarray[tuple[int, int], np.dtype[np.float32]]
+_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: ...
diff --git a/IM/ims.py b/IM/ims.py
index eded55e1..b3787da6 100644
--- a/IM/ims.py
+++ b/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)))
 
 
diff --git a/IM/konno_ohmachi.py b/IM/konno_ohmachi.py
index 18023f45..caccd7c0 100644
--- a/IM/konno_ohmachi.py
+++ b/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.
diff --git a/tests/test_core_stub.py b/tests/test_core_stub.py
new file mode 100644
index 00000000..8abf73c1
--- /dev/null
+++ b/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

From 5f3a3fe78ed78c4e7d3aa1b0ec9cc6edaa05cf50 Mon Sep 17 00:00:00 2001
From: Jake Faulkner <jake.faulkner@canterbury.ac.nz>
Date: Tue, 6 Oct 2026 08:48:43 +1300
Subject: [PATCH 2/2] ruff fixes

---
 IM/_core.pyi | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/IM/_core.pyi b/IM/_core.pyi
index e2734074..5790b9db 100644
--- a/IM/_core.pyi
+++ b/IM/_core.pyi
@@ -6,10 +6,10 @@ Every array argument must be a C-contiguous array of the stated dtype.
 
 import numpy as np
 
-_Array1F64 = np.ndarray[tuple[int], np.dtype[np.float64]]
-_Array2F64 = np.ndarray[tuple[int, int], np.dtype[np.float64]]
-_Array2F32 = np.ndarray[tuple[int, int], np.dtype[np.float32]]
-_Array3F64 = np.ndarray[tuple[int, int, int], np.dtype[np.float64]]
+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: ...