Move from pyright to ty

.pre-commit-config.yaml

@@ -8,19 +8,8 @@ repos:
       args: [ --extend-select, I, --fix ]
     # Run the formatter.
     - id: ruff-format
-- repo: https://github.com/RobertCraigie/pyright-python
-  rev: v1.1.405
+- repo: https://github.com/astral-sh/ty-pre-commit
+  # ty version.
+  rev: v0.0.50
   hooks:
-  - id: pyright
-    additional_dependencies: [jax>=0.4.35, numba>=0.60.0]
-    exclude: ^tests/
-    # Change this to automatically get from pyproject.toml or a requirements.txt?
-    # Seemingly there are good reasons why pre-commit cannot do this https://github.com/pre-commit/pre-commit/issues/730
-    # So we have to duplicate the dependencies here, but suggestions are welcome on how to improve this pyright integration!
-
-
-
-
-
-
-
+    - id: ty

cuthbert/discrete/smoother.py

@@ -53,8 +53,9 @@ def build_smoother(get_trans_matrix: GetTransitionMatrix) -> Smoother:
 
 def smoother_prepare(
     filter_state: DiscreteFilterState,
-    get_trans_matrix: GetTransitionMatrix,
     model_inputs: ArrayTreeLike,
+    *,
+    get_trans_matrix: GetTransitionMatrix,
     key: KeyArray | None = None,
 ) -> DiscreteSmootherState:
     """Prepare a state for a smoother step.

cuthbert/factorial/gaussian.py

@@ -1,3 +1,6 @@
+# ty: ignore[invalid-return-type, unknown-argument]
+# ty cannot unify NamedTuple `_replace`/`Self` with the constrained `KalmanState`
+# TypeVar; the proper fix would be @overload instead of the TypeVar.
 """Factorial utilities for Kalman states."""
 
 from typing import TypeVar
@@ -66,7 +69,7 @@ def extract(factorial_state: KalmanState, factorial_inds: ArrayLike) -> KalmanSt
     new_elem = tree.map(lambda x: _extract_arr(x, factorial_inds), factorial_state.elem)
     new_state = factorial_state._replace(elem=new_elem)
 
-    if isinstance(factorial_state, LinearizedKalmanFilterState):
+    if isinstance(new_state, LinearizedKalmanFilterState):
         new_mean_prev = _extract_arr(factorial_state.mean_prev, factorial_inds)
         new_state = new_state._replace(mean_prev=new_mean_prev)
 

cuthbert/factorial/smc.py

@@ -1,3 +1,7 @@
+# ty: ignore[invalid-return-type, unknown-argument]
+# ty cannot unify NamedTuple `_replace`/`Self` with the constrained
+# `GeneralParticleFilterState` TypeVar; the proper fix would be @overload
+# instead of the TypeVar.
 """Factorial utilities for SMC particle-filter states."""
 
 from typing import TypeVar

cuthbert/gaussian/kalman.py

@@ -211,8 +211,9 @@ def filter_combine(
 
 def smoother_prepare(
     filter_state: KalmanFilterState,
-    get_dynamics_params: GetDynamicsParams,
     model_inputs: ArrayTreeLike,
+    *,
+    get_dynamics_params: GetDynamicsParams,
     store_gain: bool = False,
     store_chol_cov_given_next: bool = False,
     key: KeyArray | None = None,

cuthbert/gaussian/moments/smoother.py

@@ -70,8 +70,9 @@ def build_smoother(
 
 def smoother_prepare(
     filter_state: LinearizedKalmanFilterState,
-    get_dynamics_params: GetDynamicsMoments,
     model_inputs: ArrayTreeLike,
+    *,
+    get_dynamics_params: GetDynamicsMoments,
     store_gain: bool = False,
     store_chol_cov_given_next: bool = False,
     key: KeyArray | None = None,

cuthbert/gaussian/taylor/smoother.py

@@ -97,8 +97,9 @@ def build_smoother(
 
 def smoother_prepare(
     filter_state: LinearizedKalmanFilterState,
-    get_dynamics_log_density: GetDynamicsLogDensity,
     model_inputs: ArrayTreeLike,
+    *,
+    get_dynamics_log_density: GetDynamicsLogDensity,
     rtol: float | None = None,
     ignore_nan_dims: bool = False,
     store_gain: bool = False,

cuthbert/inference.py

@@ -11,6 +11,7 @@ class InitPrepare(Protocol):
     def __call__(
         self,
         model_inputs: ArrayTreeLike,
+        *,
         key: KeyArray | None = None,
     ) -> ArrayTree:
         """Prepare the initial state for the inference.
@@ -34,6 +35,7 @@ class FilterPrepare(Protocol):
     def __call__(
         self,
         model_inputs: ArrayTreeLike,
+        *,
         key: KeyArray | None = None,
     ) -> ArrayTree:
         """Prepare the state for the filter at the next time point.
@@ -87,6 +89,7 @@ def __call__(
         self,
         filter_state: ArrayTreeLike,
         model_inputs: ArrayTreeLike,
+        *,
         key: KeyArray | None = None,
     ) -> ArrayTree:
         """Prepare the state for the smoother at the next time point.
@@ -155,6 +158,7 @@ def __call__(
         self,
         filter_state: ArrayTreeLike,
         model_inputs: ArrayTreeLike | None = None,
+        *,
         key: KeyArray | None = None,
     ) -> ArrayTree:
         """Convert the filter state to a smoother state.

cuthbert/smc/backward_sampler.py

@@ -74,9 +74,10 @@ def build_smoother(
 
 def convert_filter_to_smoother_state(
     filter_state: ParticleFilterState,
+    model_inputs: ArrayTreeLike | None = None,
+    *,
     resampling: Resampling,
     n_smoother_particles: int,
-    model_inputs: ArrayTreeLike | None = None,
     key: KeyArray | None = None,
 ) -> ParticleSmootherState:
     """Convert a particle filter state to a particle smoother state.

cuthbertlib/linearize/moments.py

@@ -63,14 +63,14 @@ def linearize_moments(
         - [sqrt-parallel-smoothers](https://github.com/EEA-sensors/sqrt-parallel-smoothers/blob/main/parsmooth/linearization/_extended.py)
     """
     if has_aux:
-        mean_and_chol_cov_function = cast(
+        mean_and_chol_cov_function_with_aux = cast(
             MeanAndCholCovFuncAux, mean_and_chol_cov_function
         )
 
         def mean_and_chol_cov_function_wrapper_aux(
             x: ArrayLike,
         ) -> tuple[Array, tuple[Array, Array, ArrayTree]]:
-            mean, chol_cov, aux = mean_and_chol_cov_function(x)
+            mean, chol_cov, aux = mean_and_chol_cov_function_with_aux(x)
             return mean, (mean, chol_cov, aux)
 
         F, (m, *extra) = jax.jacfwd(

cuthbertlib/resampling/protocols.py

@@ -1,3 +1,6 @@
+# ty: ignore[invalid-return-type]
+# f-strings are not supported docstrings and ty is confused about this
+# See https://github.com/state-space-models/cuthbert/issues/261
 """Shared protocols for resampling algorithms."""
 
 from typing import Protocol, runtime_checkable

pkg/cuthbert/pyproject.toml

@@ -21,7 +21,7 @@ dependencies = [
 [project.optional-dependencies]
 tests = [
     "chex",
-    "pre-commit", "ruff", "pyright",
+    "pre-commit", "ruff", "ty",
     "pytest", "pytest-xdist",
     "entangled-cli; python_version >= '3.12'",
 ]

pyproject.toml

@@ -18,6 +18,9 @@ extend-select = ["D"]
 [tool.ruff.lint.pydocstyle]
 convention = "google"
 
+[tool.ty.src]
+exclude = ["tests"]
+
 [tool.pytest]
 [tool.pytest.ini_options]
 markers = "examples: Run tangled example scripts as tests"