v4

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,20 @@ 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)
+    all_ = []
+    for n in dir(mod):
+        if not n.startswith("_") and hasattr(mod, n):
+            all_.append(n)
+            globals_[n] = getattr(mod, n)
+    return all_
+
 
+__all__ = ["get_xp", "clone_module"]
 
 def __dir__() -> list[str]:
     return __all__

array_api_compat/dask/array/__init__.py

@@ -1,12 +1,8 @@
 from typing import Final
 
-from dask.array import *  # noqa: F403
+from ..._internal import clone_module
 
-# 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
+__all__ = clone_module("dask.array", globals())
 
 # These imports may overwrite names from the import * above.
 from . import _aliases
@@ -20,6 +16,11 @@
 __import__(__package__ + '.linalg')
 __import__(__package__ + '.fft')
 
-__all__ += _aliases.__all__
-__all__ += ["__array_api_version__", "__array_namespace_info__", "linalg", "fft"]
-__all__ = sorted(set(__all__))
+__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/fft.py

@@ -1,10 +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 fft_all
-_n: dict[str, object] = {}
-exec('from dask.array.fft import *', _n)
-fft_all = list(_n)
-del _n
+from ..._internal import clone_module
+
+__all__ = clone_module("dask.array.fft", globals())
 
 from ...common import _fft
 from ..._internal import get_xp
@@ -14,7 +10,7 @@
 fftfreq = get_xp(da)(_fft.fftfreq)
 rfftfreq = get_xp(da)(_fft.rfftfreq)
 
-__all__ = fft_all + ["fftfreq", "rfftfreq"]
+__all__ += ["fftfreq", "rfftfreq"]
 
 def __dir__() -> list[str]:
     return __all__

array_api_compat/dask/array/linalg.py

@@ -8,20 +8,13 @@
 from dask.array import matmul, outer, tensordot
 
 # Exports
-from dask.array.linalg import *  # noqa: F403
-
-from ..._internal import get_xp
+from ..._internal import clone_module, get_xp
 from ...common import _linalg
 from ...common._typing import Array as _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 = {}
-exec('from dask.array.linalg import *', _n)
-linalg_all = list(_n)
-del _n
+__all__ = clone_module("dask.array.linalg", globals())
+
+from ._aliases import matrix_transpose, vecdot
 
 EighResult = _linalg.EighResult
 QRResult = _linalg.QRResult
@@ -61,11 +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"]
 
 def __dir__() -> list[str]:
     return __all__

array_api_compat/numpy/__init__.py

@@ -1,14 +1,9 @@
 # ruff: noqa: PLC0414
 from typing import Final
 
-import numpy as np
-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
+__all__ = clone_module("numpy", globals())
 
 # These imports may overwrite names from the import * above.
 from . import _aliases
@@ -31,7 +26,7 @@
 __array_api_version__: Final = "2024.12"
 
 __all__ = sorted(
-    set(np.__all__) 
+    set(__all__) 
     | set(_aliases.__all__) 
     | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"}
 )

array_api_compat/numpy/fft.py

@@ -1,8 +1,8 @@
 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 clone_module
+
+__all__ = clone_module("numpy.fft", globals())
 
 from .._internal import get_xp
 from ..common import _fft

array_api_compat/numpy/linalg.py

@@ -7,18 +7,17 @@
 
 import numpy as np
 
-from numpy.linalg import *  # noqa: F403
-
-from .._internal import get_xp
+from .._internal import clone_module, get_xp
 from ..common import _linalg
 
+from .._internal import clone_module
+
+__all__ = clone_module("numpy.linalg", globals())
+
 # These functions are in both the main and linalg namespaces
 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

array_api_compat/torch/__init__.py

@@ -1,24 +1,8 @@
 from typing import Final
 
-from torch import * # noqa: F403
+from .._internal import clone_module
 
-# Several names are not included in the above import *
-_torch_dir = set()
-import torch
-for n in dir(torch):
-    if (n.startswith('_')
-        or n.endswith('_')
-        or 'backward' in n):
-        continue
-    exec(f"{n} = torch.{n}")
-    _torch_dir.add(n)
-del n
-
-# torch.__all__ is wildly incorrect
-_n: dict[str, object] = {}
-exec('from torch import *', _n)
-_torch_all = set(_n)
-del _n
+__all__ = clone_module("torch", globals())
 
 # These imports may overwrite names from the import * above.
 from . import _aliases
@@ -32,11 +16,10 @@
 __array_api_version__: Final = '2024.12'
 
 __all__ = sorted(
-    set(_torch_all)
+    set(__all__)
     | set(_aliases.__all__)
     | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"}
-    | {"from_dlpack"}
 )
 
 def __dir__() -> list[str]:
-    return sorted(set(__all__) | set(_torch_dir))
+    return __all__

array_api_compat/torch/fft.py

@@ -4,13 +4,11 @@
 
 import torch
 import torch.fft
-from torch.fft import * # noqa: F403
 
 from ._typing import Array
+from .._internal import clone_module
 
-# 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__})
+__all__ = clone_module("torch.fft", globals())
 
 # Several torch fft functions do not map axes to dim
 
@@ -77,10 +75,7 @@ def ifftshift(
     return torch.fft.ifftshift(x, dim=axes, **kwargs)
 
 
-__all__ = sorted(
-    set(__all__)
-    | {"fftn", "ifftn", "rfftn", "irfftn", "fftshift", "ifftshift"}
-)
+__all__ += ["fftn", "ifftn", "rfftn", "irfftn", "fftshift", "ifftshift"]
 
 def __dir__() -> list[str]:
     return __all__

array_api_compat/torch/linalg.py

@@ -3,12 +3,9 @@
 import torch
 from typing import Optional, Union, Tuple
 
-from torch.linalg import * # noqa: F403
+from .._internal import clone_module
 
-# torch.linalg doesn't define __all__
-# from torch.linalg import __all__ as linalg_all
-from torch import linalg as torch_linalg
-linalg_all = [i for i in dir(torch_linalg) if not i.startswith('_')]
+__all__ = clone_module("torch.linalg", globals())
 
 # outer is implemented in torch but aren't in the linalg namespace
 from torch import outer
@@ -110,8 +107,8 @@ def vector_norm(
         return out
     return torch.linalg.vector_norm(x, ord=ord, axis=axis, keepdim=keepdims, **kwargs)
 
-__all__ = linalg_all + ['outer', 'matmul', 'matrix_transpose', 'tensordot',
-                        'cross', 'vecdot', 'solve', 'trace', 'vector_norm']
+__all__ += ['outer', 'matmul', 'matrix_transpose', 'tensordot',
+            'cross', 'vecdot', 'solve', 'trace', 'vector_norm']
 
 def __dir__() -> list[str]:
     return __all__