Handle undefined types in annotations in the new function stubs

Except for NestedSequence, which isn't a real type from typing.

Author: asmeurer

array_api_tests/function_stubs/_types.py

@@ -10,6 +10,9 @@
 array = TypeVar('array')
 device = TypeVar('device')
 dtype = TypeVar('dtype')
-
-__all__ = ['Literal', 'Optional', 'Tuple', 'Union', 'array', 'device', 'dtype']
+SupportsDLPack = TypeVar('SupportsDLPack')
+SupportsBufferProtocol = TypeVar('SupportsBufferProtocol')
+PyCapsule = TypeVar('PyCapsule')
+__all__ = ['Literal', 'Optional', 'Tuple', 'Union', 'array', 'device',
+'dtype', 'SupportsDLPack', 'SupportsBufferProtocol', 'PyCapsule']
 

array_api_tests/function_stubs/array_object.py

@@ -10,7 +10,8 @@
 
 from __future__ import annotations
 
-from ._types import Optional, Tuple, Union, array
+from enum import IntEnum
+from ._types import Optional, PyCapsule, Tuple, Union, array
 
 def __abs__(x: array, /) -> array:
     """
@@ -42,7 +43,7 @@ def __dlpack__(*, stream: Optional[int] = None) -> PyCapsule:
     """
     pass
 
-def __dlpack_device__() -> Tuple[enum.IntEnum, int]:
+def __dlpack_device__() -> Tuple[IntEnum, int]:
     """
     Note: __dlpack_device__ is a method of the array object.
     """

array_api_tests/function_stubs/creation_functions.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from ._types import Optional, Tuple, Union, array, device, dtype
+from ._types import Optional, SupportsBufferProtocol, SupportsDLPack, Tuple, Union, array, device, dtype
 
 def arange(start: Union[int, float], /, *, stop: Optional[Union[int, float]] = None, step: Union[int, float] = 1, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     pass

array_api_tests/function_stubs/data_type_functions.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 from ._types import Union, array, dtype
+from collections.abc import Sequence
 
 def finfo(type: Union[dtype, array], /) -> finfo:
     pass

generate_stubs.py

@@ -39,8 +39,10 @@
 
 from __future__ import annotations
 
+from enum import *
 from ._types import *
 from .constants import *
+from collections.abc import *
 '''
 # ^ Constants are used in some of the type annotations
 
@@ -88,8 +90,11 @@
 array = TypeVar('array')
 device = TypeVar('device')
 dtype = TypeVar('dtype')
-
-__all__ = ['Literal', 'Optional', 'Tuple', 'Union', 'array', 'device', 'dtype']
+SupportsDLPack = TypeVar('SupportsDLPack')
+SupportsBufferProtocol = TypeVar('SupportsBufferProtocol')
+PyCapsule = TypeVar('PyCapsule')
+__all__ = ['Literal', 'Optional', 'Tuple', 'Union', 'array', 'device',
+'dtype', 'SupportsDLPack', 'SupportsBufferProtocol', 'PyCapsule']
 
 '''
 def main():
@@ -645,6 +650,7 @@ def clean_type(typ):
     typ = typ.replace('\\', '')
     typ = typ.replace(' ', '')
     typ = typ.replace(',', ', ')
+    typ = typ.replace('enum.', '')
     return typ
 
 def add_annotation(sig, annotation):