v3

Author: crusaderky

array_api_compat/cupy/fft.py

@@ -30,7 +30,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,27 +1,25 @@
 from typing import Final
 
-import dask.array as da
 from dask.array import *  # noqa: F403
 
+# The above is missing a wealth of stuff
+import dask.array as da
+__all__ = [n for n in dir(da) if not n.startswith("_")]
+globals().update({n: getattr(da, n) for n in __all__})
+del da
+
 # 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"
+del Final
 
 # 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))
+__all__ += _aliases.__all__
+__all__ += ["__array_api_version__", "__array_namespace_info__", "linalg", "fft"]
+__all__ = sorted(set(__all__))

array_api_compat/numpy/fft.py

@@ -1,6 +1,9 @@
 import numpy as np
 from numpy.fft import *  # noqa: F403
 
+__all__ = [n for n in dir(np.fft) if not n.startswith("_")]
+globals().update({n: getattr(np.fft, n) for n in __all__})
+
 from .._internal import get_xp
 from ..common import _fft
 
@@ -20,7 +23,8 @@
 ifftshift = get_xp(np)(_fft.ifftshift)
 
 
-__all__ = sorted(set(np.fft.__all__) | set(_fft.__all__))
+__all__ = sorted(set(__all__) | set(_fft.__all__))
 
 def __dir__() -> list[str]:
-    return sorted(set(dir(np.fft)) | set(_fft.__all__))
+    return __all__
+

array_api_compat/numpy/linalg.py

@@ -16,6 +16,9 @@
 from ._aliases import matmul, matrix_transpose, tensordot, vecdot  # noqa: F401
 from ._typing import Array
 
+__all__ = [n for n in dir(np.linalg) if not n.startswith("_")]
+globals().update({n: getattr(np.linalg, n) for n in __all__})
+
 cross = get_xp(np)(_linalg.cross)
 outer = get_xp(np)(_linalg.outer)
 EighResult = _linalg.EighResult
@@ -122,7 +125,7 @@ def solve(x1: Array, x2: Array, /) -> Array:
     "tensorsolve",
     "vector_norm",
 ]
-__all__ = sorted(set(np.linalg.__all__) | set(_linalg.__all__) | set(_all))
+__all__ = sorted(set(__all__) | set(_linalg.__all__) | set(_all))
 
 def __dir__() -> list[str]:
-    return sorted(set(dir(np.linalg)) | set(_linalg.__all__) | set(_all))
+    return __all__

array_api_compat/torch/fft.py

@@ -8,6 +8,10 @@
 
 from ._typing import Array
 
+# The above is missing a wealth of stuff
+__all__ = [n for n in dir(torch.fft) if not n.startswith("_")]
+globals().update({n: getattr(torch.fft, n) for n in __all__})
+
 # Several torch fft functions do not map axes to dim
 
 def fftn(
@@ -73,15 +77,10 @@ def ifftshift(
     return torch.fft.ifftshift(x, dim=axes, **kwargs)
 
 
-_all = {
-    "fftn",
-    "ifftn",
-    "rfftn",
-    "irfftn",
-    "fftshift",
-    "ifftshift",
-}
-__all__ = sorted(set(torch.fft.__all__) |_all)
+__all__ = sorted(
+    set(__all__)
+    | {"fftn", "ifftn", "rfftn", "irfftn", "fftshift", "ifftshift"}
+)
 
 def __dir__() -> list[str]:
-    return sorted(set(dir(torch.fft)) | _all)
+    return __all__

tests/test_all.py

@@ -21,7 +21,7 @@
         "nan",
         "newaxis",
         "pi",
-        # Creation functions
+        # Creation Functions
         "arange",
         "asarray",
         "empty",
@@ -157,7 +157,7 @@
         "nonzero",
         "searchsorted",
         "where",
-        # Set functions
+        # Set Functions
         "unique_all",
         "unique_counts",
         "unique_inverse",
@@ -239,15 +239,13 @@
 
 
 def all_names(mod):
-    """Return all names imported by `from mod import *`.
-    This is typically `__all__` but, if not defined, Python
-    implements automated fallbacks.
-    """
-    # Note: this method also makes the test trip if a name is
-    # in __all__ but doesn't actually appear in the module.
+    """Return all names available in a module."""
     objs = {}
     exec(f"from {mod.__name__} import *", objs)
-    return list(objs)
+    for n in dir(mod):
+        if not n.startswith("_") and hasattr(mod, n):
+            objs[n] = getattr(mod, n)
+    return set(objs)
 
 
 def get_mod(library, module, *, compat):
@@ -257,49 +255,46 @@ def get_mod(library, module, *, compat):
     return getattr(xp, module) if module else xp
 
 
-@pytest.mark.parametrize("func", [all_names, dir])
 @pytest.mark.parametrize("module", list(NAMES))
 @pytest.mark.parametrize("library", wrapped_libraries)
-def test_array_api_names(library, module, func):
-    """Test that __all__ and dir() aren't missing any exports
+def test_array_api_names(library, module):
+    """Test that __all__ isn't missing any exports
     dictated by the Standard.
     """
     mod = get_mod(library, module, compat=True)
-    missing = set(NAMES[module]) - set(func(mod))
+    missing = set(NAMES[module]) - all_names(mod)
     xfail = set(XFAILS.get((library, module), []))
     xpass = xfail - missing
     fails = missing - xfail
     assert not xpass, f"Names in XFAILS are defined: {xpass}"
     assert not fails, f"Missing exports: {fails}"
 
 
-@pytest.mark.parametrize("func", [all_names, dir])
 @pytest.mark.parametrize("module", list(NAMES))
 @pytest.mark.parametrize("library", wrapped_libraries)
-def test_compat_doesnt_hide_names(library, module, func):
+def test_compat_doesnt_hide_names(library, module):
     """The base namespace can have more names than the ones explicitly exported
     by array-api-compat. Test that we're not suppressing them.
     """
     bare_mod = get_mod(library, module, compat=False)
     compat_mod = get_mod(library, module, compat=True)
 
-    missing = set(func(bare_mod)) - set(func(compat_mod))
+    missing = all_names(bare_mod) - all_names(compat_mod)
     missing = {name for name in missing if not name.startswith("_")}
     assert not missing, f"Non-Array API names have been hidden: {missing}"
 
 
-@pytest.mark.parametrize("func", [all_names, dir])
 @pytest.mark.parametrize("module", list(NAMES))
 @pytest.mark.parametrize("library", wrapped_libraries)
-def test_compat_doesnt_add_names(library, module, func):
+def test_compat_doesnt_add_names(library, module):
     """Test that array-api-compat isn't adding names to the namespace
     besides those defined by the Array API Standard.
     """
     bare_mod = get_mod(library, module, compat=False)
     compat_mod = get_mod(library, module, compat=True)
 
     aapi_names = set(NAMES[module])
-    spurious = set(func(compat_mod)) - set(func(bare_mod)) - aapi_names - {"__all__"}
+    spurious = all_names(compat_mod) - all_names(bare_mod) - aapi_names
     # Quietly ignore *Result dataclasses
     spurious = {name for name in spurious if not name.endswith("Result")}
     assert not spurious, (