TYP: fix typing errors in numpy.linalg

Author: jorenham

array_api_compat/numpy/linalg.py

@@ -1,14 +1,35 @@
-from numpy.linalg import * # noqa: F403
-from numpy.linalg import __all__ as linalg_all
-import numpy as _np
+# pyright: reportAttributeAccessIssue=false
+# pyright: reportUnknownArgumentType=false
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false
+
+from __future__ import annotations
+
+import numpy as np
+
+# intersection of `np.linalg.__all__` on numpy 1.22 and 2.2, minus `_linalg.__all__`
+from numpy.linalg import (
+    LinAlgError,
+    cond,
+    det,
+    eig,
+    eigvals,
+    eigvalsh,
+    inv,
+    lstsq,
+    matrix_power,
+    multi_dot,
+    norm,
+    tensorinv,
+    tensorsolve,
+)
 
-from ..common import _linalg
 from .._internal import get_xp
+from ..common import _linalg
 
 # These functions are in both the main and linalg namespaces
-from ._aliases import matmul, matrix_transpose, tensordot, vecdot # noqa: F401
-
-import numpy as np
+from ._aliases import matmul, matrix_transpose, tensordot, vecdot  # noqa: F401
+from ._typing import Array
 
 cross = get_xp(np)(_linalg.cross)
 outer = get_xp(np)(_linalg.outer)
@@ -38,19 +59,28 @@
 # To workaround this, the below is the code from np.linalg.solve except
 # only calling solve1 in the exactly 1D case.
 
+
 # This code is here instead of in common because it is numpy specific. Also
 # note that CuPy's solve() does not currently support broadcasting (see
 # https://github.com/cupy/cupy/blob/main/cupy/cublas.py#L43).
-def solve(x1: _np.ndarray, x2: _np.ndarray, /) -> _np.ndarray:
+def solve(x1: Array, x2: Array, /) -> Array:
     try:
         from numpy.linalg._linalg import (
-        _makearray, _assert_stacked_2d, _assert_stacked_square,
-        _commonType, isComplexType, _raise_linalgerror_singular
+            _assert_stacked_2d,
+            _assert_stacked_square,
+            _commonType,
+            _makearray,
+            _raise_linalgerror_singular,
+            isComplexType,
         )
     except ImportError:
         from numpy.linalg.linalg import (
-        _makearray, _assert_stacked_2d, _assert_stacked_square,
-        _commonType, isComplexType, _raise_linalgerror_singular
+            _assert_stacked_2d,
+            _assert_stacked_square,
+            _commonType,
+            _makearray,
+            _raise_linalgerror_singular,
+            isComplexType,
         )
     from numpy.linalg import _umath_linalg
 
@@ -61,30 +91,53 @@ def solve(x1: _np.ndarray, x2: _np.ndarray, /) -> _np.ndarray:
     t, result_t = _commonType(x1, x2)
 
     # This part is different from np.linalg.solve
+    gufunc: np.ufunc
     if x2.ndim == 1:
         gufunc = _umath_linalg.solve1
     else:
         gufunc = _umath_linalg.solve
 
     # This does nothing currently but is left in because it will be relevant
     # when complex dtype support is added to the spec in 2022.
-    signature = 'DD->D' if isComplexType(t) else 'dd->d'
-    with _np.errstate(call=_raise_linalgerror_singular, invalid='call',
-                      over='ignore', divide='ignore', under='ignore'):
-        r = gufunc(x1, x2, signature=signature)
+    signature = "DD->D" if isComplexType(t) else "dd->d"
+    with np.errstate(
+        call=_raise_linalgerror_singular,
+        invalid="call",
+        over="ignore",
+        divide="ignore",
+        under="ignore",
+    ):
+        r: Array = gufunc(x1, x2, signature=signature)
 
     return wrap(r.astype(result_t, copy=False))
 
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
-if hasattr(np.linalg, 'vector_norm'):
+if hasattr(np.linalg, "vector_norm"):
     vector_norm = np.linalg.vector_norm
 else:
     vector_norm = get_xp(np)(_linalg.vector_norm)
 
-__all__ = linalg_all + _linalg.__all__ + ['solve']
 
-del get_xp
-del np
-del linalg_all
-del _linalg
+__all__ = [
+    "LinAlgError",
+    "cond",
+    "det",
+    "eig",
+    "eigvals",
+    "eigvalsh",
+    "inv",
+    "lstsq",
+    "matrix_power",
+    "multi_dot",
+    "norm",
+    "tensorinv",
+    "tensorsolve",
+]
+__all__ += _linalg.__all__
+__all__ += ["solve", "vector_norm"]
+
+
+def __dir__() -> list[str]:
+    return __all__