Add automatically-generated .pyi files to cuda_core

.pre-commit-config.yaml

@@ -19,8 +19,9 @@ repos:
     hooks:
       - id: ruff-check
         args: [--fix, --show-fixes]
-        exclude: ^cuda_bindings/cuda/bindings/_internal/_fast_enum\.py$
+        exclude: (^cuda_bindings/cuda/bindings/_internal/_fast_enum\.py$)|(.*\.pyi$)
       - id: ruff-format
+        exclude: .*\.pyi$
 
   - repo: local
     hooks:
@@ -42,6 +43,16 @@ repos:
         language: system
         files: '^.*/docs/source/.*\.md$'
 
+      - id: stubgen-pyx-cuda-core
+        name: Generate .pyi stubs for cuda_core
+        entry: stubgen-pyx cuda_core/cuda --continue-on-error --include-private
+        language: python
+        files: ^cuda_core/cuda/.*\.(pyx|pxd)$
+        pass_filenames: false
+        additional_dependencies:
+          - stubgen-pyx==0.2.6
+          - Cython==3.2.4
+
   # Standard hooks
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: "3e8a8703264a2f4a69428a0aa4dcb512790b2c8c"  # frozen: v6.0.0
@@ -56,7 +67,7 @@ repos:
     - id: check-yaml
     - id: debug-statements
     - id: end-of-file-fixer
-      exclude: &gen_exclude '^(?:cuda_python/README\.md|cuda_bindings/cuda/bindings/.*\.in?|cuda_bindings/docs/source/module/.*\.rst?)$'
+      exclude: &gen_exclude '^(?:cuda_python/README\.md|cuda_bindings/cuda/bindings/.*\.in?|cuda_bindings/docs/source/module/.*\.rst?|.*\.pyi)$'
     - id: mixed-line-ending
     - id: trailing-whitespace
       exclude: |
@@ -79,9 +90,18 @@ repos:
     rev: 8e5c80792e2ec0c87804d8ef915bf35e2caea6da  # frozen: v1.20.0
     hooks:
       - id: mypy
+        alias: mypy-pathfinder
         name: mypy-pathfinder
         files: ^cuda_pathfinder/cuda/.*\.py$  # Exclude tests directory
         args: [--config-file=cuda_pathfinder/pyproject.toml]
+      - id: mypy
+        alias: mypy-cuda-core
+        name: mypy-cuda-core
+        files: ^cuda_core/cuda/.*\.(py|pyi)$
+        pass_filenames: false
+        args: [--config-file=cuda_core/pyproject.toml, cuda_core/cuda/core]
+        additional_dependencies:
+          - numpy
 
   - repo: https://github.com/rhysd/actionlint
     rev: "914e7df21a07ef503a81201c76d2b11c789d3fca"  # frozen: v1.7.12

.spdx-ignore

@@ -13,4 +13,7 @@ cuda_core/cuda/core/_include/dlpack.h
 cuda_core/cuda/core/_include/aoti_shim.h
 cuda_core/cuda/core/_include/aoti_shim.def
 
+# Generated by stubgen-pyx; regenerated on every commit so a header would be lost
+cuda_core/cuda/**/*.pyi
+
 qa/ctk-next.drawio.svg

cuda_core/MANIFEST.in

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-recursive-include cuda/core *.pyx *.pxd *.pxi
+recursive-include cuda/core *.pyx *.pxd *.pxi *.pyi
 recursive-include cuda/core/_cpp *.cpp *.hpp
 recursive-include cuda/core/_include *.h *.hpp
+include cuda/core/py.typed

cuda_core/cuda/core/_context.pyi

@@ -0,0 +1,87 @@
+# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_context.pyx
+
+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 Stream, StreamOptions
+
+
+class Context:
+    """CUDA context wrapper.
+
+    Context objects represent CUDA contexts and cannot be instantiated directly.
+    Use Device or Stream APIs to obtain context objects.
+    """
+
+    def close(self):
+        """Release this context wrapper's underlying CUDA handles."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        ...
+
+    @property
+    def handle(self) -> cuda.bindings.driver.CUcontext | None:
+        """Return the underlying CUcontext handle."""
+
+    @property
+    def _handle(self) -> cuda.bindings.driver.CUcontext | None:
+        ...
+
+    @property
+    def is_green(self) -> bool:
+        """True if this context was created from device resources."""
+
+    @property
+    def resources(self) -> DeviceResources:
+        """Query the hardware resources provisioned for this context.
+
+        For green contexts, returns the resources this context was created
+        with (SM partition, workqueue config). For primary contexts, returns
+        the full device resources.
+
+        Raises :class:`RuntimeError` if the context has been closed.
+        """
+
+    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
+        contexts, use :meth:`Device.create_stream` instead.
+
+        Parameters
+        ----------
+        options : :obj:`~_stream.StreamOptions`, optional
+            Customizable dataclass for stream creation options.
+
+        Returns
+        -------
+        :obj:`~_stream.Stream`
+            Newly created stream object.
+        """
+
+    def __eq__(self, other: object) -> bool:
+        ...
+
+    def __hash__(self) -> int:
+        ...
+
+    def __repr__(self) -> str:
+        ...
+
+@dataclass
+class ContextOptions:
+    """Options for context creation.
+
+    Attributes
+    ----------
+    resources : :obj:`~cuda.core.typing.DeviceResourcesType`
+        Device resources used to create a green context.
+    """
+    resources: DeviceResourcesType
+__all__ = ['Context', 'ContextOptions']
+DeviceResourcesType = Sequence[SMResource | WorkqueueResource]

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