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__