Add automatically-generated .pyi files to cuda_core (#2061)

* Add infrastructure for type-checking cuda_core

* Update all type annotations so they pass type-checking

* Add new .pyi files

* Pin Cython version

* Address feedback in the PR

* Fix tests

* Complete type annotations for public APIs

* Fix type annotations

* Fix previous merge

* Remove pre-cuda.bindings imports

* Fix isinstance check

* Add CONTRIBUTING docs

* More type annotations

* Fix build

* Fix types

Author: mdboom

.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

CONTRIBUTING.md

@@ -17,11 +17,42 @@ Thank you for your interest in contributing to CUDA Python! Based on the type of
 
 ## Table of Contents
 
-- [Pre-commit](#pre-commit)
-- [Code signing](#code-signing)
-- [Developer Certificate of Origin (DCO)](#developer-certificate-of-origin-dco)
-- [CI infrastructure overview](#ci-infrastructure-overview)
-
+- [Contributing to CUDA Python](#contributing-to-cuda-python)
+  - [Table of Contents](#table-of-contents)
+  - [Type stubs for cuda.core](#type-stubs-for-cudacore)
+  - [Pre-commit](#pre-commit)
+  - [Code signing](#code-signing)
+  - [Developer Certificate of Origin (DCO)](#developer-certificate-of-origin-dco)
+  - [CI infrastructure overview](#ci-infrastructure-overview)
+    - [CI Pipeline Flow](#ci-pipeline-flow)
+    - [Pipeline Execution Details](#pipeline-execution-details)
+    - [Branch-specific Artifact Flow](#branch-specific-artifact-flow)
+      - [Main Branch](#main-branch)
+      - [Backport Branches](#backport-branches)
+    - [Key Infrastructure Details](#key-infrastructure-details)
+  - [Code coverage](#code-coverage)
+
+
+## Type stubs for cuda.core
+
+`cuda.core` is a PEP 561-compliant package: it ships a `py.typed` marker and
+`.pyi` stub files alongside the Cython extensions.  The stubs
+are checked into the repository.
+
+**You do not need to run stubgen-pyx manually.**  A pre-commit hook
+regenerates the corresponding `.pyi` files automatically when you commit.
+The results are then also tested with `mypy`.
+
+A few things to keep in mind:
+
+- **Do not edit `.pyi` files by hand.**  They are regenerated from the Cython
+  sources on every commit that touches those sources; manual edits will be
+  overwritten.
+- **Type annotations belong in the `.pyx`/`.pxd` source.**  stubgen-pyx reads
+  Cython type annotations and docstrings to build the stubs, so keeping the
+  source well-annotated is the right way to improve stub quality.
+- **To run mypy manually (outside of pre-commit)**: `python -m mypy
+  --config-file cuda_core/pyproject.toml
 
 ## Pre-commit
 This project uses [pre-commit.ci](https://pre-commit.ci/) with GitHub Actions. All pull requests are automatically checked for pre-commit compliance, and any pre-commit failures will block merging until resolved.

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/__init__.py

@@ -5,7 +5,7 @@
 from cuda.core._version import __version__
 
 
-def _import_versioned_module():
+def _import_versioned_module() -> None:
     import importlib
 
     from cuda import bindings
@@ -28,7 +28,7 @@ def _import_versioned_module():
 del _import_versioned_module
 
 
-def _patch_rlcompleter_for_cython_properties():
+def _patch_rlcompleter_for_cython_properties() -> None:
     # TODO: This can be removed when Python 3.13 is our minimum-supported version:
     #   https://github.com/python/cpython/pull/149577
 
@@ -55,13 +55,13 @@ def _patch_rlcompleter_for_cython_properties():
     # member_descriptor types, which are what Cython uses for properties on cdef
     # classes.
     class _PatchedPropMeta(type):
-        def __instancecheck__(cls, inst):
+        def __instancecheck__(cls, inst: object) -> bool:
             return isinstance(inst, (property, GetSetDescriptorType, MemberDescriptorType))
 
     class _PatchedProperty(metaclass=_PatchedPropMeta):
         pass
 
-    rlcompleter.property = _PatchedProperty
+    rlcompleter.property = _PatchedProperty  # type: ignore[attr-defined]
 
 
 _patch_rlcompleter_for_cython_properties()

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
+
+
+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: object=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
@@ -20,9 +21,12 @@ from cuda.core._resource_handles cimport (
     as_intptr,
     as_py,
 )
-from cuda.core._stream import Stream, StreamOptions
+from cuda.core._stream import Stream
 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: object = 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