v5

Author: crusaderky

array_api_compat/_internal.py

@@ -58,12 +58,17 @@ def clone_module(mod_name: str, globals_: dict[str, object]) -> list[str]:
     Returns __all__.
     """
     mod = importlib.import_module(mod_name)
-    all_ = []
+    # 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):
-            all_.append(n)
-            globals_[n] = getattr(mod, n)
-    return all_
+            objs[n] = getattr(mod, n)
+
+    globals_.update(objs)
+    return list(objs)
 
 
 __all__ = ["get_xp", "clone_module"]

array_api_compat/numpy/__init__.py

@@ -3,6 +3,9 @@
 
 from .._internal import clone_module
 
+# 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.

tests/test_all.py

@@ -5,6 +5,8 @@
 import numpy as np
 import pytest
 
+from array_api_compat._internal import clone_module
+
 from ._helpers import wrapped_libraries
 
 NAMES = {
@@ -241,10 +243,7 @@
 def all_names(mod):
     """Return all names available in a module."""
     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)
+    clone_module(mod.__name__, objs)
     return set(objs)