address review comments

Author: lucascolley

.github/workflows/ci.yml

@@ -38,6 +38,7 @@ jobs:
         run: |
           pixi run -e lint pylint
           pixi run -e lint mypy
+          pixi run -e lint pyright
 
   checks:
     name: Check ${{ matrix.environment }}

pixi.lock

@@ -77,7 +77,7 @@ numpy = "*"
 pytest = "*"
 
 [tool.pixi.feature.lint.pypi-dependencies]
-basedmypy = "*"
+basedmypy = { version = "*", extras = ["faster-cache"] }
 basedpyright = "*"
 
 [tool.pixi.feature.lint.tasks]

pyproject.toml

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import contextlib
+import typing
 import warnings
-from typing import TYPE_CHECKING, Any
 
 # array-api-strict#6
 import array_api_strict as xp  # type: ignore[import-untyped]
@@ -12,9 +12,8 @@
 
 from array_api_extra import atleast_nd, cov, create_diagonal, expand_dims, kron, sinc
 
-if TYPE_CHECKING:
-    # To be changed to a Protocol later (see data-apis/array-api#589)
-    Array = Any  # type: ignore[no-any-explicit]
+if typing.TYPE_CHECKING:
+    from array_api_extra._typing import Array
 
 
 class TestAtLeastND:
@@ -198,7 +197,7 @@ def test_kron_smoke(self):
             ((2, 0, 0, 2), (2, 0, 2)),
         ],
     )
-    def test_kron_shape(self, shape_a: tuple[int], shape_b: tuple[int]):
+    def test_kron_shape(self, shape_a: tuple[int, ...], shape_b: tuple[int, ...]):
         a = xp.ones(shape_a)
         b = xp.ones(shape_b)
         normalised_shape_a = xp.asarray(