Rework test and imports

Author: crusaderky

array_api_compat/cupy/__init__.py

@@ -1,13 +1,24 @@
+from typing import Final
 from cupy import * # noqa: F403
 
 # from cupy import * doesn't overwrite these builtin names
 from cupy import abs, max, min, round # noqa: F401
 
 # These imports may overwrite names from the import * above.
 from ._aliases import * # noqa: F403
+from ._info import __array_namespace_info__  # noqa: F401
 
 # See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')
 __import__(__package__ + '.fft')
 
-__array_api_version__ = '2024.12'
+__array_api_version__: Final = '2024.12'
+
+__all__ = sorted(
+    {name for name in globals() if not name.startswith("__")}
+    - {"Final", "_aliases", "_info", "_typing"}
+    | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"}
+)
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/cupy/_aliases.py

@@ -7,7 +7,6 @@
 from ..common import _aliases, _helpers
 from ..common._typing import NestedSequence, SupportsBufferProtocol
 from .._internal import get_xp
-from ._info import __array_namespace_info__
 from ._typing import Array, Device, DType
 
 bool = cp.bool_
@@ -155,7 +154,7 @@ def count_nonzero(
 else:
     unstack = get_xp(cp)(_aliases.unstack)
 
-__all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'astype',
+__all__ = _aliases.__all__ + ['asarray', 'astype',
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',

array_api_compat/dask/array/__init__.py

@@ -1,12 +1,27 @@
 from typing import Final
 
+import dask.array as da
 from dask.array import *  # noqa: F403
 
 # These imports may overwrite names from the import * above.
+from . import _aliases
 from ._aliases import *  # noqa: F403
+from ._info import __array_namespace_info__  # noqa: F401
 
 __array_api_version__: Final = "2024.12"
 
 # See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')
 __import__(__package__ + '.fft')
+
+def _make_all(base):
+    return sorted(
+        set(base) 
+        | set(_aliases.__all__) 
+        | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"}
+    )
+
+__all__ = _make_all(da.__all__)
+
+def __dir__() -> list[str]:
+    return _make_all(dir(da))

array_api_compat/dask/array/_aliases.py

@@ -41,7 +41,6 @@
     NestedSequence,
     SupportsBufferProtocol,
 )
-from ._info import __array_namespace_info__
 
 isdtype = get_xp(np)(_aliases.isdtype)
 unstack = get_xp(da)(_aliases.unstack)
@@ -355,7 +354,6 @@ def count_nonzero(
 
 
 __all__ = [
-    "__array_namespace_info__",
     "count_nonzero",
     "bool",
     "int8", "int16", "int32", "int64",

array_api_compat/dask/array/fft.py

@@ -1,13 +1,10 @@
 from dask.array.fft import * # noqa: F403
 # dask.array.fft doesn't have __all__. If it is added, replace this with
-#
-# from dask.array.fft import __all__ as linalg_all
-_n = {}
+# from dask.array.fft import __all__ as fft_all
+_n: dict[str, object] = {}
 exec('from dask.array.fft import *', _n)
-for k in ("__builtins__", "Sequence", "annotations", "warnings"):
-    _n.pop(k, None)
 fft_all = list(_n)
-del _n, k
+del _n
 
 from ...common import _fft
 from ..._internal import get_xp

array_api_compat/dask/array/linalg.py

@@ -20,10 +20,8 @@
 # from dask.array.linalg import __all__ as linalg_all
 _n = {}
 exec('from dask.array.linalg import *', _n)
-for k in ('__builtins__', 'annotations', 'operator', 'warnings', 'Array'):
-    _n.pop(k, None)
 linalg_all = list(_n)
-del _n, k
+del _n
 
 EighResult = _linalg.EighResult
 QRResult = _linalg.QRResult

array_api_compat/numpy/__init__.py

@@ -1,6 +1,7 @@
 # ruff: noqa: PLC0414
 from typing import Final
 
+import numpy as np
 from numpy import *  # noqa: F403  # pyright: ignore[reportWildcardImportFromLibrary]
 
 # from numpy import * doesn't overwrite these builtin names
@@ -10,7 +11,9 @@
 from numpy import round as round
 
 # These imports may overwrite names from the import * above.
+from . import _aliases
 from ._aliases import *  # noqa: F403
+from ._info import __array_namespace_info__  # noqa: F401
 
 # Don't know why, but we have to do an absolute import to import linalg. If we
 # instead do
@@ -23,13 +26,15 @@
 
 __import__(__package__ + ".fft")
 
-from ..common._helpers import *  # noqa: F403
 from .linalg import matrix_transpose, vecdot  # noqa: F401
 
-try:
-    # Used in asarray(). Not present in older versions.
-    from numpy import _CopyMode  # noqa: F401
-except ImportError:
-    pass
-
 __array_api_version__: Final = "2024.12"
+
+__all__ = sorted(
+    set(np.__all__) 
+    | set(_aliases.__all__) 
+    | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"}
+)
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/numpy/_aliases.py

@@ -9,7 +9,6 @@
 from .._internal import get_xp
 from ..common import _aliases, _helpers
 from ..common._typing import NestedSequence, SupportsBufferProtocol
-from ._info import __array_namespace_info__
 from ._typing import Array, Device, DType
 
 if TYPE_CHECKING:
@@ -158,7 +157,6 @@ def count_nonzero(
     unstack = get_xp(np)(_aliases.unstack)
 
 __all__ = _aliases.__all__ + [
-    "__array_namespace_info__",
     "asarray",
     "astype",
     "acos",

array_api_compat/numpy/fft.py

@@ -1,5 +1,5 @@
 import numpy as np
-from numpy.fft import fft2, ifft2, irfft2, rfft2
+from numpy.fft import *  # noqa: F403
 
 from .._internal import get_xp
 from ..common import _fft
@@ -20,7 +20,7 @@
 ifftshift = get_xp(np)(_fft.ifftshift)
 
 
-__all__ = _fft.__all__ + ["fft2", "ifft2", "irfft2", "rfft2"]
+__all__ = sorted(set(np.fft.__all__) | set(_fft.__all__))
 
 def __dir__() -> list[str]:
-    return __all__
+    return sorted(set(dir(np.fft)) | set(_fft.__all__))

array_api_compat/numpy/linalg.py

@@ -7,22 +7,7 @@
 
 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 numpy.linalg import *  # noqa: F403
 
 from .._internal import get_xp
 from ..common import _linalg
@@ -120,7 +105,7 @@ def solve(x1: Array, x2: Array, /) -> Array:
     vector_norm = get_xp(np)(_linalg.vector_norm)
 
 
-__all__ = _linalg.__all__ + [
+_all = [
     "LinAlgError",
     "cond",
     "det",
@@ -137,6 +122,7 @@ def solve(x1: Array, x2: Array, /) -> Array:
     "tensorsolve",
     "vector_norm",
 ]
+__all__ = sorted(set(np.linalg.__all__) | set(_linalg.__all__) | set(_all))
 
 def __dir__() -> list[str]:
-    return __all__
+    return sorted(set(dir(np.linalg)) | set(_linalg.__all__) | set(_all))

array_api_compat/torch/__init__.py

@@ -1,6 +1,9 @@
+from typing import Final
+
 from torch import * # noqa: F403
 
 # Several names are not included in the above import *
+_torch_all = set()
 import torch
 for n in dir(torch):
     if (n.startswith('_')
@@ -10,13 +13,25 @@
         or 'backward' in n):
         continue
     exec(f"{n} = torch.{n}")
+    _torch_all.add(n)
 del n
 
 # These imports may overwrite names from the import * above.
+import _aliases
 from ._aliases import * # noqa: F403
+from ._info import __array_namespace_info__  # noqa: F401
 
 # See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')
 __import__(__package__ + '.fft')
 
-__array_api_version__ = '2024.12'
+__array_api_version__: Final = '2024.12'
+
+__all__ = sorted(
+    set(_torch_all)
+    | set(_aliases.__all__)
+    | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"}
+)
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/torch/_aliases.py

@@ -9,7 +9,6 @@
 from .._internal import get_xp
 from ..common import _aliases
 from ..common._typing import NestedSequence, SupportsBufferProtocol
-from ._info import __array_namespace_info__
 from ._typing import Array, Device, DType
 
 _int_dtypes = {
@@ -824,7 +823,7 @@ def sign(x: Array, /) -> Array:
         return out
 
 
-__all__ = ['__array_namespace_info__', 'asarray', 'result_type', 'can_cast',
+__all__ = ['asarray', 'result_type', 'can_cast',
            'permute_dims', 'bitwise_invert', 'newaxis', 'conj', 'add',
            'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
            'bitwise_right_shift', 'bitwise_xor', 'copysign', 'count_nonzero',