[FEA]: Add .pyi files so IDEs can provide autocompletion

[mdboom]

Is this a duplicate?

• [x] I confirmed there appear to be no duplicate issues for this request and that I agree to the Code of Conduct

Area

General cuda-python

Is your feature request related to a problem? Please describe.

Python language servers and type checkers (e.g. pylance, ruff-lsp, pyre, mypy) generally don't import code, so extension modules, such as those from Cython, are opaque to them.

Describe the solution you'd like

If we provided .pyi files alongside our extension modules, these tools would do a better job at type-checking and providing autocompletions etc.

While I have no direct experience with it, cythonbuilder claims to automatically generate them. Other options may include mypy's stubgen or pyright's --create-stub feature.

If that doesn't work, we may be able to generate .pyi alongside most of the generated code by modifying our generator.

~We will also need some way (through CI or the pre-commit hook, if possible) to ensure that the .pyi files are staying in-sync with the .pyx files.~

We probably don't need to check in the .pyi files into the repo if we can generate them at build time (alongside the .cpp files).

Describe alternatives you've considered

No response

Additional context

No response

[mdboom]

cythonbuilder

Seems to not support some more recent Cython syntax that we are using. I think any tool that attempts to parse Cython that /isn't/ Cython is likely to run into this maintenance burden. Similar tools Cystub and CythonPEG are in a similar state. None of these tools likes our Cython code out-of-the-box today.

There is a 2019 Cython issue to generate .pyi files from Cython.

pyright

Pyright's --createstubs command parses Python code, so isn't useful for Cython.

mypy

mypy's stubgen command gets the closest. It doesn't parse the .pyx file, but instead imports and inspects the .so module. It does a good job of extracting what's in there, but it doesn't get all the way. For example, here is a class:

class CUDA_BATCH_MEM_OP_NODE_PARAMS_v2_st:
    count: Incomplete
    ctx: Incomplete
    flags: Incomplete
    paramArray: Incomplete
    def __init__(self, *args, **kwargs) -> None: ...
    def getPtr(self) -> Any: ...
    def __reduce__(self): ...

class CUDA_CHILD_GRAPH_NODE_PARAMS(CUDA_CHILD_GRAPH_NODE_PARAMS_st):
    @classmethod
    def __init__(cls, *args, **kwargs) -> None: ...

So, it's getting the inheritance right, but the types of the members are marked as Incomplete. (I'm not sure those are introspectable).

Functions, likewise, lack a signature. This is surprising, since an introspectable signature is present -- mypy isn't picking it up, though, maybe because the cython function doesn't look enough like a regular Python function.

cuArray3DCreate: _cython_3_0_12.cython_function_or_method

While this isn't super useful for type-checking, this alone would be a useful improvement to autocompletion. For example, in VSCode:

Before:

After:

My own take is:

• mypy is the leading option for this. mypy is indisputably a "popular and maintained" package.
• It might be worth a short investigation to see if some combination of Cython build flags and changes to mypy can improve the amount of information in the .pyi file.
• It's an open question as to whether it's worth including these "partially specified" .pyi files.

EDIT: Relevant mypy issue

[mdboom]

After some time looking at this, here's my current thinking:

Stubs

Note that stubs serve two purposes: IDE autocompletion and type-checking. Many more developers use IDE autocompletion, and generally it handles "correct, but not fully specified" bindings quite well. The bar to having something useful for type-checking is much higher.

In my view, there are 3 main options to start including .pyi stubs with cuda_bindings and other Cython-using, cybind-generated bindings:

Option 1: Extend Cython to output .pyi files

In a perfect world, Cython would be extended to generate .pyi files alongside the .cpp files, since that is the layer that has all of the information necessary. However, making that work is a pretty significant undertaking, largely because of the historical divergence of typing in Cython and Python, and I don't think there is necessarily buy-in from the core Cython developers (who have suggested "Option 3" below). There is an issue from 2019 that gets into some of the details. There are a number of external tools that generate .pyi from .pyd and .pyx files but none of them seem to keep up with all Cython features. Relying on any of them would probably make it harder to upgrade to new versions of Cython for new features:

• CyStub: uses cython compiler to generate PYX -> PYI
• cythonbuilder: custom parser PYX -> PYI
• CythonPeg: pyparsing custom definitions PYX -> PYI

Option 2: Generate .pyi files from cython_gen and cybind

Our own generator tools could just generate .pyi files alongside the .pyx and .pxd files. An upside is that we would be in full control over exactly what goes into the .pyi files.

The downside to this approach is that it will only work with the generated code. Handwritten Cython code would need to have handwritten .pyi files alongside it, and use a tool such as stubtest to confirm that the two remain in sync. We might be able to use GenAI to create the first revision of these non-generated stub files with some work, but we would need a way to segment the work to avoid token limits. (Hilariously, when I tried, Copilot got "bored" half way through and said "just keep doing the same here".)

Option 3: Use tools that introspect the compiled modules to generate .pyi files

So that leaves general Python tools that generate stubs. mypy's stubgen tool is the only suitable choice for Cython code, as it can import a compiled extension module and inspect it.

Ideally, .pyi file generation would just happen transparently as part of the build. Once the .so files have been compiled, we should be able to run stubgen over them and include the .pyi files in the wheel. (We can also include stub files in another wheel, if the additional disk space is deemed too onerous). This would require a bit of build tool wizardry to solve, but let's assume it can be, as I think that approach is preferable over anything that would include them in the repository. A side effect of this is that we could detect when the public API of a Cython module has changed, by detecting when the generated .pyi file (ignoring docstrings) changes.

I have spent the most time experimenting with this option. Some notes on details to fix below:

Including Cython functions

stubgen doesn't currently recognize Cython functions as a Python callable. There was some work on this back in 2020 that stalled out, but things seem to have become considerably easier since then. This simple patch to mypy seems to be all that is required. We can submit upstream (the hard part will be writing a good test):

     def is_function(self, obj: object) -> bool:
-        if self.is_c_module:
-            return inspect.isbuiltin(obj)
-        else:
-            return inspect.isfunction(obj)
+        return inspect.isroutine(obj)

Cython properties

Properties are a weird corner of the Python typing standard. While the fget and fset functions in a property have signatures, those don't seem be used by mypy, instead it parses a Google/Numpy format docstring on the property to determine its type. Additionally, Cython properties are a different implementation and don't have an accessible fget function anyway.

It would be nice to specify the type this way (in the return type of the getter):

    @property
    def Width(self) -> int:
        return 0
    @Width.setter
    def Width(self, int Width):
        ...

This causes Cython to generate a docstring for the property of the form:

{name}: {type}

Unfortunately, this is at odds with the Google/Numpy property docstring format mypy expects, which is:

{type}: {description}

So, instead, we can specify the type like this in a docstring:

    @property
    def Width(self):
        "int: Width of the thing"
        return 0
    @Width.setter
    def Width(self, int Width):
        ...

We can use this second approach today, and it's probably fine since we generate the bindings anyway. Or we could do some combination of (a) fixing Cython's generated docstring to match the Google/Numpy format or (b) helping mypy understand Cython's format. Both options may be hard given that they are at least partially-breaking changes.

Include private bug

mypy stubgen currently includes private (_-prefixed) things in C extension modules, even when you tell it not too. This is a known bug which seems like it should be easy to fix at first glance.

Overloads

For some functions, mypy stubgen will generate false "overloaded" versions because it is confused by the docstring syntax.

For example, the docstring for cuCtxCreate is (abridged):

cuCtxCreate(ctxCreateParams : Optional[CUctxCreateParams], unsigned int flags, dev)
Create a CUDA context.

Creates a new CUDA context and associates it with the calling thread.
The `flags` parameter is described below. The context is created with a
usage count of 1 and the caller of :py:obj:`~.cuCtxCreate()` must call
:py:obj:`~.cuCtxDestroy()` when done using the context. If a context is
already current to the thread, it is supplanted by the newly created
context and may be restored by a subsequent call to
:py:obj:`~.cuCtxPopCurrent()`.

stubgen first sees the correct signature at the beginning of the docstring (with 3 arguments), but also later sees the :py:obj:`~.cuCtxCreate()` reference and thinks that is another signature with zero arguments. It then generates this in the stub:

@overload
def cuCtxCreate(ctxCreateParams: CUctxCreateParams | None, flags: int, dev) -> Any: ...
@overload
def cuCtxCreate() -> Any: ...

Some combination of improving mypy here or what we generate in the docstring could probably address this.

[mdboom]

Somewhat related to this discussion is to decide whether to continue documenting the C++ types in Python docstrings, rather than the Python types that they map to.

[Andy-Jost]

Issue #1442 is a duplicate of this issue. Below is a comment from @mdboom, reproduced here to keep the discussion in one place. I will close the other ticket.

I agree that hand-written .pyi files would be a huge burden.

I see a couple ways forward:

• Use a tool that can import an already built Python extension and generate stubs from it, such as stubgen in "inspect mode". That tools doesn't do a perfect job because our docstrings (particularly the generated ones) use a mixture of C and Python types. I go into more detail about that in https://github.com/NVIDIA/cuda-python/issues/928#issuecomment-3266879676
• Convert our Cython files to use the new "Pure python" syntax, and then use a stubgen in "no import mode". This is a simple process -- it basically just strips the implementation from all functions and outputs the result.

In either path, you integrate stubgen as part of the build process, and you wouldn't check the .pyi files into source control.

Both of these paths are significant work, though I feel like the second one is ultimately more sustainable from a "leveraging where the Python community is going" perspective. We could do it incrementally -- convert a single .pyx file to the new syntax, and learn how difficult that is, whether LLMs or other tooling are helpful, etc.

[Andy-Jost]

Convert our Cython files to use the new "Pure python" syntax,

My initial impression of this approach is that it would be quite onerous, at least for cuda‑core (perhaps less so for generated code). We would need to rewrite all .pyx files and switch to writing Cython in an alternative style that seems more verbose and harder to maintain.

Here’s an illustrative example from the link you provided:

# Cython (a declaration similar to this appears in a .pxd)
cdef int func_not_needing_the_gil() nogil:
    return 1

versus

# Python (appears in a .py)
@cython.nogil
@cython.cfunc
@cython.returns(cython.int)
def func_not_needing_the_gil() -> cython.int:
    return 1

The pure Python syntax is much longer and, at a glance, more difficult to parse. It also forces developers to keep two very different styles in sync: the concise, native Cython declarations we use today and a decorator‑heavy pure‑Python representation. Matching these up correctly requires specialized knowledge and quite a bit of mental parsing. I'm worried it would make everyday development and code review more cumbersome.

[mdboom]

My understanding is that the "pure-Python" syntax approach was added to Cython to avoid all of the problems that having a separate, incompatible syntax bring to Cython. This is mainly that none of the tooling that works on pure Python (linters, type checkers, formatters etc.) can work with it, and all of these tools have to be rewritten specially the Cython ecosystem. The ones that exist are all pretty poor imitations given the much smaller user base. Cython long predates type annotations being added to Python, and I'm certain if starting today the ad hoc Cython syntax (or I guess it was pyrex) would not have been invented.

The "pure Python" syntax is certainly more verbose, but I think it is a matter of personal preference. I have trouble remembering all the ways that Cython syntax deviates from Python syntax, so I would welcome bringing them closer. And, IME, the ability to load a Cython module as pure Python and get a full-fledged debugging experience is worth ALL of the downsides.

It also forces developers to keep two very different styles in sync: the concise, native Cython declarations we use today and a decorator‑heavy pure‑Python representation.

During the transition period, of course it will be easy to confuse the two as we will be working with both styles in the same code base for a while. But I'm not sure what you mean by keeping in sync. There would still only be one copy of the code...

EDIT: Ah, I get you are talking about the .py vs. .pxd files and keeping those in sync. Indeed that's an issue and it's a shame Cython doesn't have an answer there.