feat: ARRAY_API_TESTS_MODULE for runtime-defined xp

Author: lucascolley

array_api_tests/__init__.py

@@ -13,19 +13,27 @@
 # You can comment the following out and instead import the specific array module
 # you want to test, e.g. `import array_api_strict 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)
+    env_var = os.environ["ARRAY_API_TESTS_MODULE"]
+    if env_var.startswith("exec(") and env_var.endswith(")"):
+        script = env_var[5:][:-1]
+        namespace = {}
+        exec(script, namespace)
+        xp = namespace["xp"]
+        xp_name = xp.__name__
+    else:
+        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 "