Make test_signatures.py use new stubs

Author: honno

array_api_tests/pytest_helpers.py

@@ -64,18 +64,8 @@ def doesnt_raise(function, message=""):
         raise AssertionError(f"Unexpected exception {e!r}")
 
 
-all_funcs = []
-for funcs in [
-    stubs.array_methods,
-    *list(stubs.category_to_funcs.values()),
-    *list(stubs.extension_to_funcs.values()),
-]:
-    all_funcs.extend(funcs)
-name_to_func = {f.__name__: f for f in all_funcs}
-
-
 def nargs(func_name):
-    return len(getfullargspec(name_to_func[func_name]).args)
+    return len(getfullargspec(stubs.name_to_func[func_name]).args)
 
 
 def fmt_kw(kw: Dict[str, Any]) -> str:

array_api_tests/stubs.py

@@ -1,12 +1,18 @@
-import sys
 import inspect
+import sys
 from importlib import import_module
 from importlib.util import find_spec
 from pathlib import Path
 from types import FunctionType, ModuleType
 from typing import Dict, List
 
-__all__ = ["array_methods", "category_to_funcs", "EXTENSIONS", "extension_to_funcs"]
+__all__ = [
+    "name_to_func",
+    "array_methods",
+    "category_to_funcs",
+    "EXTENSIONS",
+    "extension_to_funcs",
+]
 
 
 spec_dir = Path(__file__).parent.parent / "array-api" / "spec" / "API_specification"
@@ -29,7 +35,6 @@
     if n != "__init__"  # probably exists for Sphinx
 ]
 
-
 category_to_funcs: Dict[str, List[FunctionType]] = {}
 for name, mod in name_to_mod.items():
     if name.endswith("_functions"):
@@ -45,3 +50,8 @@
     objects = [getattr(mod, name) for name in mod.__all__]
     assert all(isinstance(o, FunctionType) for o in objects)
     extension_to_funcs[ext] = objects
+
+all_funcs = []
+for funcs in [array_methods, *category_to_funcs.values(), *extension_to_funcs.values()]:
+    all_funcs.extend(funcs)
+name_to_func: Dict[str, FunctionType] = {f.__name__: f for f in all_funcs}

array_api_tests/test_signatures.py

@@ -1,46 +1,44 @@
 import inspect
+from itertools import chain
 
 import pytest
 
 from ._array_module import mod, mod_name, ones, eye, float64, bool, int64, _UndefinedStub
 from .pytest_helpers import raises, doesnt_raise
 from . import dtype_helpers as dh
 
-from . import function_stubs
 from . import stubs
 
 
-def stub_module(name):
-    for m in stubs.extensions:
-        if name in getattr(function_stubs, m).__all__:
-            return m
+def extension_module(name) -> bool:
+    for funcs in stubs.extension_to_funcs.values():
+        for func in funcs:
+            if name == func.__name__:
+                return True
+    else:
+        return False
 
-def extension_module(name):
-    return name in stubs.extensions and name in function_stubs.__all__
 
-extension_module_names = []
-for n in function_stubs.__all__:
-    if extension_module(n):
-        extension_module_names.extend([f'{n}.{i}' for i in getattr(function_stubs, n).__all__])
+params = []
+for name in [f.__name__ for funcs in stubs.category_to_funcs.values() for f in funcs]:
+    if name in ["where", "expand_dims", "reshape"]:
+        params.append(pytest.param(name, marks=pytest.mark.skip(reason="faulty test")))
+    else:
+        params.append(name)
 
 
-params = []
-for name in function_stubs.__all__:
-    marks = []
-    if extension_module(name):
-        marks.append(pytest.mark.xp_extension(name))
-    params.append(pytest.param(name, marks=marks))
-for name in extension_module_names:
-    ext = name.split('.')[0]
-    mark = pytest.mark.xp_extension(ext)
-    params.append(pytest.param(name, marks=[mark]))
+for ext, name in [(ext, f.__name__) for ext, funcs in stubs.extension_to_funcs.items() for f in funcs]:
+    params.append(pytest.param(name, marks=pytest.mark.xp_extension(ext)))
 
 
-def array_method(name):
-    return stub_module(name) == 'array_object'
+def array_method(name) -> bool:
+    return name in [f.__name__ for f in stubs.array_methods]
 
-def function_category(name):
-    return stub_module(name).rsplit('_', 1)[0].replace('_', ' ')
+def function_category(name) -> str:
+    for category, funcs in chain(stubs.category_to_funcs.items(), stubs.extension_to_funcs.items()):
+        for func in funcs:
+            if name == func.__name__:
+                return category
 
 def example_argument(arg, func_name, dtype):
     """
@@ -138,7 +136,7 @@ def example_argument(arg, func_name, dtype):
             return ones((3,), dtype=dtype)
         # Linear algebra functions tend to error if the input isn't "nice" as
         # a matrix
-        elif arg.startswith('x') and func_name in function_stubs.linalg.__all__:
+        elif arg.startswith('x') and func_name in [f.__name__ for f in stubs.extension_to_funcs["linalg"]]:
             return eye(3)
         return known_args[arg]
     else:
@@ -147,13 +145,15 @@ def example_argument(arg, func_name, dtype):
 @pytest.mark.parametrize('name', params)
 def test_has_names(name):
     if extension_module(name):
-        assert hasattr(mod, name), f'{mod_name} is missing the {name} extension'
-    elif '.' in name:
-        extension_mod, name = name.split('.')
-        assert hasattr(getattr(mod, extension_mod), name), f"{mod_name} is missing the {function_category(name)} extension function {name}()"
+        ext = next(
+            ext for ext, funcs in stubs.extension_to_funcs.items()
+            if name in [f.__name__ for f in funcs]
+        )
+        ext_mod = getattr(mod, ext)
+        assert hasattr(ext_mod, name), f"{mod_name} is missing the {function_category(name)} extension function {name}()"
     elif array_method(name):
         arr = ones((1, 1))
-        if getattr(function_stubs.array_object, name) is None:
+        if name not in [f.__name__ for f in stubs.array_methods]:
             assert hasattr(arr, name), f"The array object is missing the attribute {name}"
         else:
             assert hasattr(arr, name), f"The array object is missing the method {name}()"
@@ -192,14 +192,12 @@ def test_function_positional_args(name):
             _mod = ones((), dtype=float64)
         else:
             _mod = example_argument('self', name, dtype)
-        stub_func = getattr(function_stubs, name)
     elif '.' in name:
         extension_module_name, name = name.split('.')
         _mod = getattr(mod, extension_module_name)
-        stub_func = getattr(getattr(function_stubs, extension_module_name), name)
     else:
         _mod = mod
-        stub_func = getattr(function_stubs, name)
+    stub_func = stubs.name_to_func[name]
 
     if not hasattr(_mod, name):
         pytest.skip(f"{mod_name} does not have {name}(), skipping.")
@@ -245,14 +243,12 @@ def test_function_keyword_only_args(name):
 
     if array_method(name):
         _mod = ones((1, 1))
-        stub_func = getattr(function_stubs, name)
     elif '.' in name:
         extension_module_name, name = name.split('.')
         _mod = getattr(mod, extension_module_name)
-        stub_func = getattr(getattr(function_stubs, extension_module_name), name)
     else:
         _mod = mod
-        stub_func = getattr(function_stubs, name)
+    stub_func = stubs.name_to_func[name]
 
     if not hasattr(_mod, name):
         pytest.skip(f"{mod_name} does not have {name}(), skipping.")