tests: fix all pyright complaints in tests

Author: rg936672

.pylintrc

@@ -438,6 +438,8 @@ disable=too-many-lines,
 # Throws false positives that prevent abstract class hierarchies; The check is
 # handled at instantiation time by other linting rules.
         abstract-method,
+# Ruff's line-too-long handling is better - no need to duplicate it here
+        line-too-long,
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

tests/coverage/compare.py

@@ -121,7 +121,7 @@ def get_most_recent_coverage_total(reference_directory: Path) -> float:
         print("**WARNING: No historic coverage data found.**")
         return 0
 
-    most_recent_file = max(files.keys(), key=files.get)
+    most_recent_file = max(files.keys(), key=files.__getitem__)
 
     with open(most_recent_file, "r", encoding="utf8") as f:
         coverage_dict = json.load(f)

tests/performance/cases/basic_coresets.py

@@ -56,6 +56,7 @@ def _setup_dataset():
         random_state=random_seed,
         return_centers=True,
     )
+    x = jnp.asarray(x)
 
     # Setup the original data object
     data = Data(x)
@@ -92,7 +93,7 @@ def setup_stein():
     data, length_scale = _setup_dataset()
     # We use kernel density matching rather than sliced score matching as it's much
     # faster than the sliced score matching used in the original unit test
-    matcher = KernelDensityMatching(length_scale=length_scale)
+    matcher = KernelDensityMatching(length_scale=length_scale.item())
     stein_kernel = SteinKernel(
         PCIMQKernel(length_scale=length_scale),
         matcher.match(jnp.asarray(data)),

tests/performance/compare.py

@@ -161,7 +161,7 @@ def get_most_recent_historic_data(
         # fine since there aren't any results present to normalise
         return {"results": {}, "normalisation": {"compilation": 0.0, "execution": 0.0}}
 
-    most_recent_file = max(files.keys(), key=files.get)
+    most_recent_file = max(files.keys(), key=files.__getitem__)
 
     with open(most_recent_file, "r", encoding="utf8") as f:
         return json.load(f)

tests/unit/test_approximation.py

@@ -26,7 +26,7 @@
 import jax.random as jr
 import numpy as np
 import pytest
-from jax.typing import ArrayLike
+from jax import Array
 
 from coreax.approximation import (
     ANNchorApproximateKernel,
@@ -42,11 +42,11 @@
 
 class _Problem(NamedTuple):
     random_key: KeyArrayLike
-    data: ArrayLike
+    data: Array
     kernel: ScalarValuedKernel
     num_kernel_points: int
     num_train_points: int
-    true_distances: ArrayLike
+    true_distances: Array
 
 
 @pytest.mark.parametrize(
@@ -113,7 +113,7 @@ def problem(self) -> _Problem:
 
         # We can repeat the above, but changing the point with which we are comparing
         # to get:
-        true_distances = np.array(
+        true_distances = jnp.array(
             [
                 0.5855723855138795,
                 0.5737865795122914,

tests/unit/test_benchmark.py

@@ -24,7 +24,7 @@
 import pytest
 import torch
 from jax import random
-from torch.utils.data import Dataset
+from torchvision.datasets import VisionDataset
 
 from benchmark.mnist_benchmark import (
     MLP,
@@ -43,10 +43,11 @@
 )
 
 
-class MockDataset(Dataset):
+class MockDataset(VisionDataset):
     """Mock dataset class for testing purposes."""
 
-    def __init__(self, data: torch.Tensor, labels: torch.Tensor) -> None:
+    # We deliberately don't call super().__init__(), as this is a mock class
+    def __init__(self, data: torch.Tensor, labels: torch.Tensor) -> None:  # pylint: disable=super-init-not-called
         """
         Initialise the MockDataset.
 

tests/unit/test_coreset.py

@@ -109,12 +109,12 @@ def test_deprecated_nodes(
         with pytest.warns(DeprecationWarning):
             nodes = coreset.nodes
 
-        if setup.coreset_type is PseudoCoreset:
+        if isinstance(coreset, PseudoCoreset):
             assert nodes == coreset.points
-        elif setup.coreset_type is Coresubset:
+        elif isinstance(coreset, Coresubset):
             assert nodes == coreset.indices
         else:
-            raise ValueError(setup.coreset_type)
+            raise TypeError(type(coreset))
 
     def test_deprecated_coreset(
         self,
@@ -189,7 +189,8 @@ def test_build_array_conversion(
             )
             with ctx:
                 coreset_from_arrays = setup.coreset_type(
-                    coreset_input_final, pre_coreset_data_final
+                    coreset_input_final,  # pyright: ignore[reportArgumentType]
+                    pre_coreset_data_final,  # pyright: ignore[reportArgumentType]
                 )
         coreset_from_data = setup.coreset_type.build(
             setup.coreset_input, setup.pre_coreset_data
@@ -267,7 +268,7 @@ def test_invalid_pre_coreset_data_type(
         with pytest.raises(TypeError, match="pre_coreset_data"):
             with warnings.catch_warnings():
                 warnings.simplefilter(action="ignore", category=DeprecationWarning)
-                coreset_type.build(indices_or_nodes, object())
+                coreset_type.build(indices_or_nodes, object())  # pyright: ignore[reportArgumentType, reportCallIssue]
 
     @pytest.mark.parametrize("coreset_type", [Coresubset, PseudoCoreset])
     def test_invalid_indices_or_points_type(
@@ -280,7 +281,7 @@ def test_invalid_indices_or_points_type(
         ):
             with warnings.catch_warnings():
                 warnings.simplefilter(action="ignore", category=DeprecationWarning)
-                coreset_type(object(), pre_coreset_data)
+                coreset_type(object(), pre_coreset_data)  # pyright: ignore[reportArgumentType, reportCallIssue]
 
 
 class TestCoresubset:

tests/unit/test_kernels.py

@@ -220,19 +220,19 @@ def test_gradients(
 
     @abstractmethod
     def expected_grad_x(
-        self, x: ArrayLike, y: ArrayLike, kernel: _ScalarValuedKernel
+        self, x: Array, y: Array, kernel: _ScalarValuedKernel
     ) -> Union[Array, np.ndarray]:
         """Compute expected gradient of the kernel w.r.t ``x``."""
 
     @abstractmethod
     def expected_grad_y(
-        self, x: ArrayLike, y: ArrayLike, kernel: _ScalarValuedKernel
+        self, x: Array, y: Array, kernel: _ScalarValuedKernel
     ) -> Union[Array, np.ndarray]:
         """Compute expected gradient of the kernel w.r.t ``y``."""
 
     @abstractmethod
     def expected_divergence_x_grad_y(
-        self, x: ArrayLike, y: ArrayLike, kernel: _ScalarValuedKernel
+        self, x: Array, y: Array, kernel: _ScalarValuedKernel
     ) -> Union[Array, np.ndarray]:
         """Compute expected divergence of the kernel w.r.t ``x`` gradient ``y``."""
 
@@ -402,9 +402,7 @@ def problem(self, request: pytest.FixtureRequest, kernel: PowerKernel) -> _Probl
         )
         return _Problem(x, y, expected_distances, modified_kernel)
 
-    def expected_grad_x(
-        self, x: ArrayLike, y: ArrayLike, kernel: PowerKernel
-    ) -> np.ndarray:
+    def expected_grad_x(self, x: Array, y: Array, kernel: PowerKernel) -> np.ndarray:
         num_points, dimension = np.atleast_2d(x).shape
 
         expected_grad = (
@@ -419,9 +417,7 @@ def expected_grad_x(
 
         return np.array(expected_grad)
 
-    def expected_grad_y(
-        self, x: ArrayLike, y: ArrayLike, kernel: PowerKernel
-    ) -> np.ndarray:
+    def expected_grad_y(self, x: Array, y: Array, kernel: PowerKernel) -> np.ndarray:
         num_points, dimension = np.atleast_2d(x).shape
 
         expected_grad = (
@@ -437,7 +433,7 @@ def expected_grad_y(
         return np.array(expected_grad)
 
     def expected_divergence_x_grad_y(
-        self, x: ArrayLike, y: ArrayLike, kernel: PowerKernel
+        self, x: Array, y: Array, kernel: PowerKernel
     ) -> np.ndarray:
         divergence = self.power * (
             (
@@ -579,9 +575,7 @@ def problem(
         )
         return _Problem(x, y, expected_distances, modified_kernel)
 
-    def expected_grad_x(
-        self, x: ArrayLike, y: ArrayLike, kernel: AdditiveKernel
-    ) -> np.ndarray:
+    def expected_grad_x(self, x: Array, y: Array, kernel: AdditiveKernel) -> np.ndarray:
         num_points, dimension = np.atleast_2d(x).shape
 
         # Variable rename allows for nicer automatic formatting
@@ -595,9 +589,7 @@ def expected_grad_x(
 
         return np.array(expected_grad)
 
-    def expected_grad_y(
-        self, x: ArrayLike, y: ArrayLike, kernel: AdditiveKernel
-    ) -> np.ndarray:
+    def expected_grad_y(self, x: Array, y: Array, kernel: AdditiveKernel) -> np.ndarray:
         num_points, dimension = np.atleast_2d(x).shape
 
         # Variable rename allows for nicer automatic formatting
@@ -612,7 +604,7 @@ def expected_grad_y(
         return np.array(expected_grad)
 
     def expected_divergence_x_grad_y(
-        self, x: ArrayLike, y: ArrayLike, kernel: AdditiveKernel
+        self, x: Array, y: Array, kernel: AdditiveKernel
     ) -> np.ndarray:
         num_points, _ = np.atleast_2d(x).shape
 
@@ -691,9 +683,7 @@ def problem(
         )
         return _Problem(x, y, expected_distances, modified_kernel)
 
-    def expected_grad_x(
-        self, x: ArrayLike, y: ArrayLike, kernel: ProductKernel
-    ) -> np.ndarray:
+    def expected_grad_x(self, x: Array, y: Array, kernel: ProductKernel) -> np.ndarray:
         num_points, dimension = np.atleast_2d(x).shape
 
         # Variable rename allows for nicer automatic formatting
@@ -709,9 +699,7 @@ def expected_grad_x(
 
         return np.array(expected_grad)
 
-    def expected_grad_y(
-        self, x: ArrayLike, y: ArrayLike, kernel: ProductKernel
-    ) -> np.ndarray:
+    def expected_grad_y(self, x: Array, y: Array, kernel: ProductKernel) -> np.ndarray:
         num_points, dimension = np.atleast_2d(x).shape
 
         # Variable rename allows for nicer automatic formatting
@@ -728,7 +716,7 @@ def expected_grad_y(
         return np.array(expected_grad)
 
     def expected_divergence_x_grad_y(
-        self, x: ArrayLike, y: ArrayLike, kernel: ProductKernel
+        self, x: Array, y: Array, kernel: ProductKernel
     ) -> np.ndarray:
         # Variable rename allows for nicer automatic formatting
         k1, k2 = kernel.first_kernel, kernel.second_kernel
@@ -748,7 +736,7 @@ def test_symmetric_product_kernel(self):
         We consider a product kernel with equal input kernels and check that
         the second kernel is never called.
         """
-        x = np.array([1])
+        x = jnp.array([1])
 
         # Form two simple mocked kernels and force any == operation to return True
         first_kernel = MagicMock(spec=ScalarValuedKernel)
@@ -1362,8 +1350,10 @@ def problem(  # noqa: C901
             expected_distances = np.zeros((num_points, num_points))
             for x_idx, x_ in enumerate(x):
                 for y_idx, y_ in enumerate(y):
-                    expected_distances[x_idx, y_idx] = scipy_norm(y_, length_scale).pdf(
-                        x_
+                    expected_distances[x_idx, y_idx] = (
+                        scipy_norm(y_, length_scale)
+                        # Ignore Pyright here - the .pdf() function definitely exists!
+                        .pdf(x_)  # pyright: ignore[reportAttributeAccessIssue]
                     )
             x, y = x.reshape(-1, 1), y.reshape(-1, 1)
         elif mode == "negative_length_scale":

tests/unit/test_metrics.py

@@ -206,6 +206,8 @@ def test_mmd_random_data(
         kernel_mm = kernel.compute(y.data, y.data)
         kernel_nm = kernel.compute(x.data, y.data)
         if mode == "weighted":
+            assert isinstance(x.weights, Array)
+            assert isinstance(y.weights, Array)
             weights_nn = x.weights[..., None] * x.weights[None, ...]
             weights_mm = y.weights[..., None] * y.weights[None, ...]
             weights_nm = x.weights[..., None] * y.weights[None, ...]
@@ -355,6 +357,7 @@ def test_ksd_random_data(
         # Compute each term in the KSD formula to obtain an expected KSD.
         kernel_mm = kernel.compute(y.data, y.data)
         if mode == "weighted":
+            assert isinstance(y.weights, Array)
             weights_mm = y.weights[..., None] * y.weights[None, ...]
             expected_ksd = jnp.sqrt(jnp.average(kernel_mm, weights=weights_mm))
             output = metric.compute(x, y, laplace_correct=False, regularise=False)