Complete type annotations for public APIs

Author: mdboom

cuda_core/cuda/core/_context.pyi

@@ -5,9 +5,10 @@ from __future__ import annotations
 from collections.abc import Sequence
 from dataclasses import dataclass
 
+import cuda.bindings.driver
 from cuda.core._device_resources import (DeviceResources, SMResource,
                                          WorkqueueResource)
-from cuda.core._stream import StreamOptions
+from cuda.core._stream import Stream, StreamOptions
 
 
 class Context:
@@ -20,15 +21,15 @@ class Context:
     def close(self):
         """Release this context wrapper's underlying CUDA handles."""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         ...
 
     @property
-    def handle(self):
+    def handle(self) -> cuda.bindings.driver.CUcontext | None:
         """Return the underlying CUcontext handle."""
 
     @property
-    def _handle(self):
+    def _handle(self) -> cuda.bindings.driver.CUcontext | None:
         ...
 
     @property
@@ -46,7 +47,7 @@ class Context:
         Raises :class:`RuntimeError` if the context has been closed.
         """
 
-    def create_stream(self, options: StreamOptions | None=None):
+    def create_stream(self, options: StreamOptions | None=None) -> Stream:
         """Create a new stream bound to this green context.
 
         This method is only available on green contexts. For primary
@@ -63,7 +64,7 @@ class Context:
             Newly created stream object.
         """
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         ...
 
     def __hash__(self) -> int:

cuda_core/cuda/core/_context.pyx

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 from cuda.bindings cimport cydriver
 from cuda.core._device_resources cimport DeviceResources, SMResource, WorkqueueResource
@@ -23,6 +24,9 @@ from cuda.core._resource_handles cimport (
 from cuda.core._stream import Stream, StreamOptions
 from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
 
+if TYPE_CHECKING:
+    import cuda.bindings.driver  # no-cython-lint
+
 
 __all__ = ['Context', 'ContextOptions']
 
@@ -37,7 +41,7 @@ cdef class Context:
     Use Device or Stream APIs to obtain context objects.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         raise RuntimeError("Context objects cannot be instantiated directly. Please use Device or Stream APIs.")
 
     @staticmethod
@@ -58,7 +62,7 @@ cdef class Context:
         return Context._from_handle(cls, h_context, device_id)
 
     @property
-    def handle(self):
+    def handle(self) -> cuda.bindings.driver.CUcontext | None:
         """Return the underlying CUcontext handle."""
         if not self._h_context:
             return None
@@ -67,7 +71,7 @@ cdef class Context:
         return as_py(self._h_context)
 
     @property
-    def _handle(self):
+    def _handle(self) -> cuda.bindings.driver.CUcontext | None:
         return self.handle
 
     @property
@@ -91,7 +95,7 @@ cdef class Context:
             raise RuntimeError("Cannot query resources on a closed context")
         return DeviceResources._init_from_ctx(self._h_context, self._device_id)
 
-    def create_stream(self, options: StreamOptions | None = None):
+    def create_stream(self, options: StreamOptions | None = None) -> Stream:
         """Create a new stream bound to this green context.
 
         This method is only available on green contexts. For primary
@@ -130,7 +134,7 @@ cdef class Context:
                 )
         self._h_context.reset()
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if not isinstance(other, Context):
             return NotImplemented
         cdef Context _other = <Context>other

cuda_core/cuda/core/_device.pyi

@@ -21,11 +21,11 @@ class DeviceProperties:
     Attributes are read-only and provide information about the device.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         ...
 
     @classmethod
-    def _init(cls, handle):
+    def _init(cls, handle: int) -> DeviceProperties:
         ...
 
     @property
@@ -634,14 +634,14 @@ class Device:
     """
     __slots__ = ('_device_id', '_memory_resource', '_has_inited', '_properties', '_resources', '_uuid', '_context', '__weakref__')
 
-    def __new__(cls, device_id: Device | int | None=None):
+    def __new__(cls, device_id: Device | int | None=None) -> Device:
         ...
 
     def _check_context_initialized(self):
         ...
 
     @classmethod
-    def get_all_devices(cls):
+    def get_all_devices(cls) -> tuple[Device, ...]:
         """
         Query the available device instances.
 
@@ -737,7 +737,7 @@ class Device:
         """Return :obj:`~_memory.MemoryResource` associated with this device."""
 
     @memory_resource.setter
-    def memory_resource(self, mr):
+    def memory_resource(self, mr: MemoryResource) -> None:
         ...
 
     @property
@@ -752,19 +752,19 @@ class Device:
 
         """
 
-    def __int__(self):
+    def __int__(self) -> int:
         """Return device_id."""
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         ...
 
     def __hash__(self) -> int:
         ...
 
-    def __eq__(self, other) -> bool:
+    def __eq__(self, other: object) -> bool:
         ...
 
-    def __reduce__(self):
+    def __reduce__(self) -> tuple:
         ...
 
     def set_current(self, ctx: Context | None=None) -> Context | None:
@@ -865,7 +865,7 @@ class Device:
 
         """
 
-    def allocate(self, size, *, stream: Stream | GraphBuilder) -> Buffer:
+    def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> Buffer:
         """Allocate device memory from a specified stream.
 
         Allocates device memory of `size` bytes on the specified `stream`
@@ -891,7 +891,7 @@ class Device:
 
         """
 
-    def sync(self):
+    def sync(self) -> None:
         """Synchronize the device.
 
         Note

cuda_core/cuda/core/_device.pyx

@@ -62,11 +62,11 @@ cdef class DeviceProperties:
         int _handle
         dict _cache
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         raise RuntimeError("DeviceProperties cannot be instantiated directly. Please use Device APIs.")
 
     @classmethod
-    def _init(cls, handle):
+    def _init(cls, handle: int) -> DeviceProperties:
         cdef DeviceProperties self = DeviceProperties.__new__(cls)
         self._handle = handle
         self._cache = {}
@@ -976,7 +976,7 @@ class Device:
         "__weakref__",
     )
 
-    def __new__(cls, device_id: Device | int | None = None):
+    def __new__(cls, device_id: Device | int | None = None) -> Device:
         if isinstance(device_id, Device):
             return device_id
 
@@ -997,7 +997,7 @@ class Device:
 
 
     @classmethod
-    def get_all_devices(cls):
+    def get_all_devices(cls) -> tuple[Device, ...]:
         """
         Query the available device instances.
 
@@ -1178,7 +1178,7 @@ class Device:
         return self._memory_resource
 
     @memory_resource.setter
-    def memory_resource(self, mr):
+    def memory_resource(self, mr: MemoryResource) -> None:
         from cuda.core._memory import MemoryResource
         assert_type(mr, MemoryResource)
         self._memory_resource = mr
@@ -1196,22 +1196,22 @@ class Device:
         """
         return default_stream()
 
-    def __int__(self):
+    def __int__(self) -> int:
         """Return device_id."""
         return self._device_id
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<Device {self._device_id} ({self.name})>"
 
     def __hash__(self) -> int:
         return hash(self.uuid)
 
-    def __eq__(self, other) -> bool:
+    def __eq__(self, other: object) -> bool:
         if not isinstance(other, Device):
             return NotImplemented
         return self._device_id == other._device_id
 
-    def __reduce__(self):
+    def __reduce__(self) -> tuple:
         return Device, (self.device_id,)
 
     def set_current(self, ctx: Context | None = None) -> Context | None:
@@ -1400,7 +1400,7 @@ class Device:
         cdef Context ctx = self._context
         return cyEvent._init(cyEvent, self._device_id, ctx._h_context, options, True)
 
-    def allocate(self, size, *, stream: Stream | GraphBuilder) -> Buffer:
+    def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> Buffer:
         """Allocate device memory from a specified stream.
 
         Allocates device memory of `size` bytes on the specified `stream`
@@ -1428,7 +1428,7 @@ class Device:
         self._check_context_initialized()
         return self.memory_resource.allocate(size, stream=stream)
 
-    def sync(self):
+    def sync(self) -> None:
         """Synchronize the device.
 
         Note

cuda_core/cuda/core/_device_resources.pyi

@@ -80,7 +80,7 @@ class SMResource:
     def flags(self) -> int:
         """Raw flags from the underlying SM resource."""
 
-    def split(self, options, *, dry_run: bool=False):
+    def split(self, options: SMResourceOptions, *, dry_run: bool=False) -> tuple[list[SMResource], SMResource]:
         """Split this SM resource into groups and a remainder.
 
         Parameters
@@ -107,14 +107,14 @@ class WorkqueueResource:
     cannot be instantiated directly.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         ...
 
     @property
     def handle(self) -> int:
         """Return the address of the underlying config ``CUdevResource`` struct."""
 
-    def configure(self, options):
+    def configure(self, options: WorkqueueResourceOptions) -> None:
         """Configure the workqueue resource in place.
 
         Parameters
@@ -134,7 +134,7 @@ class DeviceResources:
     This class cannot be instantiated directly.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         ...
 
     @property

cuda_core/cuda/core/_device_resources.pyx

@@ -471,7 +471,12 @@ cdef class SMResource:
         """Raw flags from the underlying SM resource."""
         return self._flags
 
-    def split(self, options not None, *, bint dry_run=False):
+    def split(
+        self,
+        options: SMResourceOptions,
+        *,
+        bint dry_run=False
+    ) -> tuple[list[SMResource], SMResource]:
         """Split this SM resource into groups and a remainder.
 
         Parameters
@@ -508,7 +513,7 @@ cdef class WorkqueueResource:
     cannot be instantiated directly.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         raise RuntimeError(
             "WorkqueueResource cannot be instantiated directly. "
             "Use dev.resources.workqueue."
@@ -529,7 +534,7 @@ cdef class WorkqueueResource:
         """Return the address of the underlying config ``CUdevResource`` struct."""
         return <intptr_t>(&self._wq_config_resource)
 
-    def configure(self, options not None):
+    def configure(self, options: WorkqueueResourceOptions) -> None:
         """Configure the workqueue resource in place.
 
         Parameters
@@ -576,7 +581,7 @@ cdef class DeviceResources:
     This class cannot be instantiated directly.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         raise RuntimeError(
             "DeviceResources cannot be instantiated directly. "
             "Use dev.resources or ctx.resources."

cuda_core/cuda/core/_dlpack.pyi

@@ -16,7 +16,7 @@ class DLDeviceType(IntEnum):
 def make_py_capsule(buf: object, versioned: bool) -> object:
     ...
 
-def classify_dl_device(buf) -> tuple[int, int]:
+def classify_dl_device(buf: object) -> tuple[int, int]:
     """Classify a buffer into a DLPack (device_type, device_id) pair.
 
     ``buf`` must expose ``is_device_accessible``, ``is_host_accessible``,

cuda_core/cuda/core/_dlpack.pyx

@@ -89,7 +89,7 @@ cdef inline int setup_dl_tensor_layout(DLTensor* dl_tensor, object buf) except -
     return 0
 
 
-def classify_dl_device(buf) -> tuple[int, int]:
+def classify_dl_device(buf: object) -> tuple[int, int]:
     """Classify a buffer into a DLPack (device_type, device_id) pair.
 
     ``buf`` must expose ``is_device_accessible``, ``is_host_accessible``,