Define/assign xp in __init__.py

Prevents circular imports, generally makes more sense

Author: honno

README.md

@@ -36,13 +36,7 @@ You need to specify the array library to test. It can be specified via the
 $ export ARRAY_API_TESTS_MODULE=numpy.array_api
 ```
 
-Alternately, change the `array_module` variable in `array_api_tests/_array_module.py`
-line, e.g.
-
-```diff
-- array_module = None
-+ import numpy.array_api as array_module
-```
+Alternately, import/define the `xp` variable in `array_api_tests/__init__.py`.
 
 ### Run the suite
 

array_api_tests/__init__.py

@@ -1,13 +1,36 @@
+import os
 from functools import wraps
-from os import getenv
+from importlib import import_module
 
 from hypothesis import strategies as st
 from hypothesis.extra import array_api
 
 from . import _version
-from ._array_module import mod as _xp
 
-__all__ = ["api_version", "xps"]
+__all__ = ["xp", "api_version", "xps"]
+
+
+# You can comment the following out and instead import the specific array module
+# you want to test, e.g. `import numpy.array_api as xp`.
+if "ARRAY_API_TESTS_MODULE" in os.environ:
+    xp_name = os.environ["ARRAY_API_TESTS_MODULE"]
+    _module, _sub = xp_name, None
+    if "." in xp_name:
+        _module, _sub = xp_name.split(".", 1)
+    xp = import_module(_module)
+    if _sub:
+        try:
+            xp = getattr(xp, _sub)
+        except AttributeError:
+            # _sub may be a submodule that needs to be imported. WE can't
+            # do this in every case because some array modules are not
+            # submodules that can be imported (like mxnet.nd).
+            xp = import_module(xp_name)
+else:
+    raise RuntimeError(
+        "No array module specified - either edit __init__.py or set the "
+        "ARRAY_API_TESTS_MODULE environment variable."
+    )
 
 
 # We monkey patch floats() to always disable subnormals as they are out-of-scope
@@ -43,9 +66,9 @@ def _from_dtype(*a, **kw):
     pass
 
 
-api_version = getenv(
-    "ARRAY_API_TESTS_VERSION", getattr(_xp, "__array_api_version__", "2021.12")
+api_version = os.getenv(
+    "ARRAY_API_TESTS_VERSION", getattr(xp, "__array_api_version__", "2021.12")
 )
-xps = array_api.make_strategies_namespace(_xp, api_version=api_version)
+xps = array_api.make_strategies_namespace(xp, api_version=api_version)
 
 __version__ = _version.get_versions()["version"]

array_api_tests/_array_module.py

@@ -1,35 +1,5 @@
-import os
-from importlib import import_module
+from . import stubs, xp
 
-from . import stubs
-
-# Replace this with a specific array module to test it, for example,
-#
-# import numpy as array_module
-array_module = None
-
-if array_module is None:
-    if 'ARRAY_API_TESTS_MODULE' in os.environ:
-        mod_name = os.environ['ARRAY_API_TESTS_MODULE']
-        _module, _sub = mod_name, None
-        if '.' in mod_name:
-            _module, _sub = mod_name.split('.', 1)
-        mod = import_module(_module)
-        if _sub:
-            try:
-                mod = getattr(mod, _sub)
-            except AttributeError:
-                # _sub may be a submodule that needs to be imported. WE can't
-                # do this in every case because some array modules are not
-                # submodules that can be imported (like mxnet.nd).
-                mod = import_module(mod_name)
-    else:
-        raise RuntimeError("No array module specified. Either edit _array_module.py or set the ARRAY_API_TESTS_MODULE environment variable")
-else:
-    mod = array_module
-    mod_name = mod.__name__
-# Names from the spec. This is what should actually be imported from this
-# file.
 
 class _UndefinedStub:
     """
@@ -45,7 +15,7 @@ def __init__(self, name):
         self.name = name
 
     def _raise(self, *args, **kwargs):
-        raise AssertionError(f"{self.name} is not defined in {mod_name}")
+        raise AssertionError(f"{self.name} is not defined in {xp.__name__}")
 
     def __repr__(self):
         return f"<undefined stub for {self.name!r}>"
@@ -67,6 +37,6 @@ def __repr__(self):
 
 for attr in _top_level_attrs:
     try:
-        globals()[attr] = getattr(mod, attr)
+        globals()[attr] = getattr(xp, attr)
     except AttributeError:
         globals()[attr] = _UndefinedStub(attr)

array_api_tests/dtype_helpers.py

@@ -6,7 +6,7 @@
 from warnings import warn
 
 from . import api_version
-from ._array_module import mod as xp
+from . import xp
 from .stubs import name_to_func
 from .typing import DataType, ScalarType
 

array_api_tests/stubs.py

@@ -6,6 +6,8 @@
 from types import FunctionType, ModuleType
 from typing import Dict, List
 
+from . import api_version
+
 __all__ = [
     "name_to_func",
     "array_methods",
@@ -15,10 +17,9 @@
     "extension_to_funcs",
 ]
 
-spec_version = "2022.12"
-spec_module = "_" + spec_version.replace('.', '_')
+spec_module = "_" + api_version.replace('.', '_')
 
-spec_dir = Path(__file__).parent.parent / "array-api" / "spec" / spec_version / "API_specification"
+spec_dir = Path(__file__).parent.parent / "array-api" / "spec" / api_version / "API_specification"
 assert spec_dir.exists(), f"{spec_dir} not found - try `git submodule update --init`"
 sigs_dir = Path(__file__).parent.parent / "array-api" / "src" / "array_api_stubs" / spec_module
 assert sigs_dir.exists()

array_api_tests/test_array_object.py

@@ -13,7 +13,7 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from ._array_module import mod as _xp
+from . import xp as _xp
 from .typing import DataType, Index, Param, Scalar, ScalarType, Shape
 
 pytestmark = pytest.mark.ci

array_api_tests/test_constants.py

@@ -4,7 +4,7 @@
 import pytest
 
 from . import dtype_helpers as dh
-from ._array_module import mod as xp
+from . import xp
 from .typing import Array
 
 pytestmark = pytest.mark.ci

array_api_tests/test_data_type_functions.py

@@ -11,7 +11,7 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from ._array_module import mod as _xp
+from . import xp as _xp
 from .typing import DataType
 
 pytestmark = pytest.mark.ci

array_api_tests/test_fft.py

@@ -14,7 +14,7 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from ._array_module import mod as xp
+from . import xp
 
 pytestmark = [
     pytest.mark.ci,

array_api_tests/test_has_names.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from ._array_module import mod as xp, mod_name
+from . import xp
 from .stubs import (array_attributes, array_methods, category_to_funcs,
                     extension_to_funcs, EXTENSIONS)
 
@@ -27,13 +27,13 @@
 def test_has_names(category, name):
     if category in EXTENSIONS:
         ext_mod = getattr(xp, category)
-        assert hasattr(ext_mod, name), f"{mod_name} is missing the {category} extension function {name}()"
+        assert hasattr(ext_mod, name), f"{xp.__name__} is missing the {category} extension function {name}()"
     elif category.startswith('array_'):
         # TODO: This would fail if ones() is missing.
         arr = xp.ones((1, 1))
         if category == 'array_attribute':
-            assert hasattr(arr, name), f"The {mod_name} array object is missing the attribute {name}"
+            assert hasattr(arr, name), f"The {xp.__name__} array object is missing the attribute {name}"
         else:
-            assert hasattr(arr, name), f"The {mod_name} array object is missing the method {name}()"
+            assert hasattr(arr, name), f"The {xp.__name__} array object is missing the method {name}()"
     else:
-        assert hasattr(xp, name), f"{mod_name} is missing the {category} function {name}()"
+        assert hasattr(xp, name), f"{xp.__name__} is missing the {category} function {name}()"

array_api_tests/test_signatures.py

@@ -30,7 +30,7 @@ def squeeze(x, /, axis):
 import pytest
 
 from . import dtype_helpers as dh
-from ._array_module import mod as xp
+from . import xp
 from .stubs import array_methods, category_to_funcs, extension_to_funcs, name_to_func
 
 pytestmark = pytest.mark.ci

array_api_tests/test_special_cases.py

@@ -33,7 +33,7 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from ._array_module import mod as xp
+from . import xp
 from .stubs import category_to_funcs
 
 pytestmark = pytest.mark.ci

reporting.py

@@ -49,7 +49,7 @@ def pytest_metadata(metadata):
     """
     Additional global metadata for --json-report.
     """
-    metadata['array_api_tests_module'] = xp.mod_name
+    metadata['array_api_tests_module'] = xp.__name__
     metadata['array_api_tests_version'] = __version__
 
 @fixture(autouse=True)