ENH: Review exported symbols; redesign test_all (#315)

Review and discussion at #315

Author: crusaderky

array_api_compat/_internal.py

@@ -2,6 +2,7 @@
 Internal helpers
 """
 
+import importlib
 from collections.abc import Callable
 from functools import wraps
 from inspect import signature
@@ -52,8 +53,25 @@ def wrapped_f(*args: object, **kwargs: object) -> object:
     return inner
 
 
-__all__ = ["get_xp"]
+def clone_module(mod_name: str, globals_: dict[str, object]) -> list[str]:
+    """Import everything from module, updating globals().
+    Returns __all__.
+    """
+    mod = importlib.import_module(mod_name)
+    # Neither of these two methods is sufficient by itself,
+    # depending on various idiosyncrasies of the libraries we're wrapping.
+    objs = {}
+    exec(f"from {mod.__name__} import *", objs)
+
+    for n in dir(mod):
+        if not n.startswith("_") and hasattr(mod, n):
+            objs[n] = getattr(mod, n)
+
+    globals_.update(objs)
+    return list(objs)
+
 
+__all__ = ["get_xp", "clone_module"]
 
 def __dir__() -> list[str]:
     return __all__

array_api_compat/common/_aliases.py

@@ -721,8 +721,6 @@ def iinfo(type_: DType | Array, /, xp: Namespace) -> Any:
     "finfo",
     "iinfo",
 ]
-_all_ignore = ["inspect", "array_namespace", "NamedTuple"]
-
 
 def __dir__() -> list[str]:
     return __all__

array_api_compat/common/_helpers.py

@@ -1062,7 +1062,5 @@ def is_lazy_array(x: object) -> TypeGuard[_ArrayApiObj]:
     "to_device",
 ]
 
-_all_ignore = ['lru_cache', 'sys', 'math', 'inspect', 'warnings']
-
 def __dir__() -> list[str]:
     return __all__

array_api_compat/common/_linalg.py

@@ -225,8 +225,6 @@ def trace(
            'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal',
            'trace']
 
-_all_ignore = ['math', 'normalize_axis_tuple', 'get_xp', 'np', 'isdtype']
-
 
 def __dir__() -> list[str]:
     return __all__

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_
@@ -141,7 +140,7 @@ def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
 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/cupy/_typing.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 __all__ = ["Array", "DType", "Device"]
-_all_ignore = ["cp"]
 
 from typing import TYPE_CHECKING
 

array_api_compat/cupy/fft.py

@@ -31,7 +31,6 @@
 
 __all__ = fft_all + _fft.__all__
 
-del get_xp
-del cp
-del fft_all
-del _fft
+def __dir__() -> list[str]:
+    return __all__
+

array_api_compat/cupy/linalg.py

@@ -43,7 +43,5 @@
 
 __all__ = linalg_all + _linalg.__all__
 
-del get_xp
-del cp
-del linalg_all
-del _linalg
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/dask/array/__init__.py

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

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",
@@ -369,8 +367,6 @@ def count_nonzero(
     "bitwise_left_shift", "bitwise_right_shift", "bitwise_invert",
 ]  # fmt: skip
 __all__ += _aliases.__all__
-_all_ignore = ["array_namespace", "get_xp", "da", "np"]
-
 
 def __dir__() -> list[str]:
     return __all__

array_api_compat/dask/array/fft.py

@@ -1,13 +1,6 @@
-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: 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
+from ..._internal import clone_module
+
+__all__ = clone_module("dask.array.fft", globals())
 
 from ...common import _fft
 from ..._internal import get_xp
@@ -17,5 +10,7 @@
 fftfreq = get_xp(da)(_fft.fftfreq)
 rfftfreq = get_xp(da)(_fft.rfftfreq)
 
-__all__ = fft_all + ["fftfreq", "rfftfreq"]
-_all_ignore = ["da", "fft_all", "get_xp", "warnings"]
+__all__ += ["fftfreq", "rfftfreq"]
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/dask/array/linalg.py

@@ -4,27 +4,17 @@
 
 import dask.array as da
 
-# Exports
-from dask.array.linalg import *  # noqa: F403
-from dask.array import outer
 # The `matmul` and `tensordot` functions are in both the main and linalg namespaces
-from dask.array import matmul, tensordot
-
+from dask.array import matmul, outer, tensordot
 
-from ..._internal import get_xp
+# Exports
+from ..._internal import clone_module, get_xp
 from ...common import _linalg
 from ...common._typing import Array
-from ._aliases import matrix_transpose, vecdot
 
-# dask.array.linalg doesn't have __all__. If it is added, replace this with
-#
-# from dask.array.linalg import __all__ as linalg_all
-_n: dict[str, object] = {}
-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
+__all__ = clone_module("dask.array.linalg", globals())
+
+from ._aliases import matrix_transpose, vecdot
 
 EighResult = _linalg.EighResult
 QRResult = _linalg.QRResult
@@ -64,10 +54,11 @@ def svdvals(x: Array) -> Array:
 vector_norm = get_xp(da)(_linalg.vector_norm)
 diagonal = get_xp(da)(_linalg.diagonal)
 
-__all__ = linalg_all + ["trace", "outer", "matmul", "tensordot",
-                        "matrix_transpose", "vecdot", "EighResult",
-                        "QRResult", "SlogdetResult", "SVDResult", "qr",
-                        "cholesky", "matrix_rank", "matrix_norm", "svdvals",
-                        "vector_norm", "diagonal"]
+__all__ += ["trace", "outer", "matmul", "tensordot",
+            "matrix_transpose", "vecdot", "EighResult",
+            "QRResult", "SlogdetResult", "SVDResult", "qr",
+            "cholesky", "matrix_rank", "matrix_norm", "svdvals",
+            "vector_norm", "diagonal"]
 
-_all_ignore = ['get_xp', 'da', 'linalg_all', 'warnings']
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/numpy/__init__.py

@@ -1,16 +1,17 @@
 # ruff: noqa: PLC0414
 from typing import Final
 
-from numpy import *  # noqa: F403  # pyright: ignore[reportWildcardImportFromLibrary]
+from .._internal import clone_module
 
-# from numpy import * doesn't overwrite these builtin names
-from numpy import abs as abs
-from numpy import max as max
-from numpy import min as min
-from numpy import round as round
+# This needs to be loaded explicitly before cloning
+import numpy.typing  # noqa: F401
+
+__all__ = clone_module("numpy", globals())
 
 # These imports may overwrite names from the import * above.
+from . import _aliases
 from ._aliases import *  # type: ignore[assignment,no-redef] # 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
@@ -26,3 +27,12 @@
 from .linalg import matrix_transpose, vecdot  # type: ignore[no-redef]  # noqa: F401
 
 __array_api_version__: Final = "2024.12"
+
+__all__ = sorted(
+    set(__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
 
 bool = np.bool_
@@ -147,8 +146,7 @@ def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
 else:
     unstack = get_xp(np)(_aliases.unstack)
 
-__all__ = [
-    "__array_namespace_info__",
+__all__ = _aliases.__all__ + [
     "asarray",
     "astype",
     "acos",
@@ -167,8 +165,6 @@ def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
     "pow",
     "take_along_axis"
 ]
-__all__ += _aliases.__all__
-_all_ignore = ["np", "get_xp"]
 
 
 def __dir__() -> list[str]:

array_api_compat/numpy/_typing.py

@@ -23,7 +23,6 @@
     Array: TypeAlias = np.ndarray
 
 __all__ = ["Array", "DType", "Device"]
-_all_ignore = ["np"]
 
 
 def __dir__() -> list[str]:

array_api_compat/numpy/fft.py

@@ -1,6 +1,8 @@
 import numpy as np
-from numpy.fft import __all__ as fft_all
-from numpy.fft import fft2, ifft2, irfft2, rfft2
+
+from .._internal import clone_module
+
+__all__ = clone_module("numpy.fft", globals())
 
 from .._internal import get_xp
 from ..common import _fft
@@ -21,15 +23,8 @@
 ifftshift = get_xp(np)(_fft.ifftshift)
 
 
-__all__ = ["rfft2", "irfft2", "fft2", "ifft2"]
-__all__ += _fft.__all__
-
+__all__ = sorted(set(__all__) | set(_fft.__all__))
 
 def __dir__() -> list[str]:
     return __all__
 
-
-del get_xp
-del np
-del fft_all
-del _fft