Update generate_stubs.py to generate stubs for the linalg extension

Author: asmeurer

array_api_tests/function_stubs/__init__.py

@@ -56,3 +56,5 @@
 from .utility_functions import all, any
 
 __all__ += ['all', 'any']
+
+from . import linalg

array_api_tests/function_stubs/linalg.py

@@ -0,0 +1,92 @@
+"""
+Function stubs for linear algebra functions (Extension).
+
+NOTE: This file is generated automatically by the generate_stubs.py script. Do
+not modify it directly.
+
+See
+https://github.com/data-apis/array-api/blob/master/spec/API_specification/linear_algebra_functions.md
+"""
+
+from __future__ import annotations
+
+
+def cholesky(x, /, *, upper=False):
+    pass
+
+def cross(x1, x2, /, *, axis=-1):
+    pass
+
+def det(x, /):
+    pass
+
+def diagonal(x, /, *, axis1=0, axis2=1, offset=0):
+    pass
+
+def eig():
+    pass
+
+def eigh(x, /, *, upper=False):
+    pass
+
+def eigvals():
+    pass
+
+def eigvalsh(x, /, *, upper=False):
+    pass
+
+def einsum():
+    pass
+
+def inv(x, /):
+    pass
+
+def lstsq(x1, x2, /, *, rtol=None):
+    pass
+
+def matmul(x1, x2, /):
+    pass
+
+def matrix_power(x, n, /):
+    pass
+
+def matrix_rank(x, /, *, rtol=None):
+    pass
+
+def norm(x, /, *, axis=None, keepdims=False, ord=None):
+    pass
+
+def outer(x1, x2, /):
+    pass
+
+def pinv(x, /, *, rtol=None):
+    pass
+
+def qr(x, /, *, mode='reduced'):
+    pass
+
+def slogdet(x, /):
+    pass
+
+def solve(x1, x2, /):
+    pass
+
+def svd(x, /, *, full_matrices=True):
+    pass
+
+def tensordot(x1, x2, /, *, axes=2):
+    pass
+
+def svdvals(x, /):
+    pass
+
+def trace(x, /, *, axis1=0, axis2=1, offset=0):
+    pass
+
+def transpose(x, /, *, axes=None):
+    pass
+
+def vecdot(x1, x2, /, *, axis=None):
+    pass
+
+__all__ = ['cholesky', 'cross', 'det', 'diagonal', 'eig', 'eigh', 'eigvals', 'eigvalsh', 'einsum', 'inv', 'lstsq', 'matmul', 'matrix_power', 'matrix_rank', 'norm', 'outer', 'pinv', 'qr', 'slogdet', 'solve', 'svd', 'tensordot', 'svdvals', 'trace', 'transpose', 'vecdot']

generate_stubs.py

@@ -116,10 +116,14 @@ def main():
             f.write(TYPES_HEADER)
 
     spec_dir = os.path.join(args.array_api_repo, 'spec', 'API_specification')
+    extensions_dir = os.path.join(args.array_api_repo, 'spec', 'extensions')
+    files = sorted([os.path.join(spec_dir, f) for f in os.listdir(spec_dir)]
+                   + [os.path.join(extensions_dir, f) for f in os.listdir(extensions_dir)])
     modules = {}
-    for filename in sorted(os.listdir(spec_dir)):
-        with open(os.path.join(spec_dir, filename)) as f:
-                  text = f.read()
+    for file in files:
+        filename = os.path.basename(file)
+        with open(file) as f:
+            text = f.read()
         functions = FUNCTION_RE.findall(text)
         methods = METHOD_RE.findall(text)
         constants = CONSTANT_RE.findall(text)
@@ -130,9 +134,18 @@ def main():
             print(f"Found signatures in {filename}")
         if not args.write:
             continue
-        py_file = filename.replace('.md', '.py')
-        py_path = os.path.join('array_api_tests', 'function_stubs', py_file)
+
         title = filename.replace('.md', '').replace('_', ' ')
+        if 'extensions' in file:
+            if filename == 'index.md':
+                continue
+            elif filename != 'linear_algebra_functions.md':
+                raise RuntimeError(f"Don't know how to handle extension file {filename}")
+            py_file = 'linalg.py'
+            title += " (Extension)"
+        else:
+            py_file = filename.replace('.md', '.py')
+        py_path = os.path.join('array_api_tests', 'function_stubs', py_file)
         module_name = py_file.replace('.py', '')
         modules[module_name] = []
         if not args.quiet:
@@ -153,6 +166,11 @@ def main():
             ismethod = sig in methods
             sig = sig.replace(r'\_', '_')
             func_name = NAME_RE.match(sig).group(1)
+            if '.' in func_name:
+                mod, func_name = func_name.split('.', 2)
+                if mod != 'linalg':
+                    raise RuntimeError(f"Unexpected namespace prefix {mod!r}")
+                sig = sig.replace(mod + '.', '')
             doc = ""
             if ismethod:
                 doc = f'''
@@ -236,6 +254,9 @@ def {annotated_sig}:{doc}
         with open(init_path, 'w') as f:
             f.write(INIT_HEADER)
             for module_name in modules:
+                if module_name == 'linalg':
+                    f.write(f'\nfrom . import {module_name}\n')
+                    continue
                 f.write(f"\nfrom .{module_name} import ")
                 f.write(', '.join(modules[module_name]))
                 f.write('\n\n')