Fix generate_stubs.py to correctly find annotations for functions with underscores

Author: asmeurer

array_api_tests/function_stubs/linalg.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 from ._types import Literal, Optional, Tuple, Union, array
+from .constants import inf
 from collections.abc import Sequence
 
 def cholesky(x: array, /, *, upper: bool = False) -> array:
@@ -46,7 +47,7 @@ def inv(x: array, /) -> array:
 def matmul(x1: array, x2: array, /) -> array:
     pass
 
-def matrix_norm(x, /, *, axis=(-2, -1), keepdims=False, ord='fro'):
+def matrix_norm(x: array, /, *, axis: Tuple[int, int] = (-2, -1), keepdims: bool = False, ord: Optional[Union[int, float, Literal[inf, -inf, 'fro', 'nuc']]] = 'fro') -> array:
     pass
 
 def matrix_power(x: array, n: int, /) -> array:
@@ -55,7 +56,7 @@ def matrix_power(x: array, n: int, /) -> array:
 def matrix_rank(x: array, /, *, rtol: Optional[Union[float, array]] = None) -> array:
     pass
 
-def matrix_transpose(x, /):
+def matrix_transpose(x: array, /) -> array:
     pass
 
 def outer(x1: array, x2: array, /) -> array:
@@ -88,7 +89,7 @@ def trace(x: array, /, *, offset: int = 0) -> array:
 def vecdot(x1: array, x2: array, /, *, axis: Optional[int] = None) -> array:
     pass
 
-def vector_norm(x, /, *, axis=None, keepdims=False, ord=2):
+def vector_norm(x: array, /, *, axis: Optional[Union[int, Tuple[int, int]]] = None, keepdims: bool = False, ord: Optional[Union[int, float, Literal[inf, -inf]]] = 2) -> array:
     pass
 
 __all__ = ['cholesky', 'cross', 'det', 'diagonal', 'eig', 'eigh', 'eigvals', 'eigvalsh', 'einsum', 'inv', 'matmul', 'matrix_norm', 'matrix_power', 'matrix_rank', 'matrix_transpose', 'outer', 'pinv', 'qr', 'slogdet', 'solve', 'svd', 'svdvals', 'tensordot', 'trace', 'vecdot', 'vector_norm']

array_api_tests/function_stubs/linear_algebra_functions.py

@@ -19,7 +19,7 @@ def einsum():
 def matmul(x1: array, x2: array, /) -> array:
     pass
 
-def matrix_transpose(x, /):
+def matrix_transpose(x: array, /) -> array:
     pass
 
 def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> array:

generate_stubs.py

@@ -690,15 +690,15 @@ def parse_annotations(spec_text, all_annotations, verbose=False):
     for line in spec_text.splitlines():
         m = HEADER_RE.match(line)
         if m:
-            name = m.group(1)
+            name = m.group(1).replace('-', '_')
             continue
         m = ALIAS_RE.match(line)
         if m:
-            alias_name = m.group(1)
+            alias_name = m.group(1).replace('-', '_')
             if alias_name not in all_annotations:
                 print(f"Warning: No annotations for aliased function {name}")
             else:
-                annotations[name] = all_annotations[m.group(1)]
+                annotations[name] = all_annotations[alias_name]
             continue
         if line == '#### Parameters':
             in_block = True