Tracking issue: stub generation

[wjakob]

Dear all (cc @cansik @torokati44 @qnzhou @tmsrise @rzhikharevich-wmt @njroussel @Speierers),

I am interested in providing a stub generation mechanism as part of nanobind. This is a tracking issue to brainstorm solutions.

Context: @cansik's nanobind-stubgen package is the only solution at the moment and works well in many cases. My goal is to overcome limitations discussed in discussion #163:

1. Enabling a better "out-of-the-box" experience by integrating stub generation into the CMake build system.

2. Stub generation currently involves complicated parsing, which is fragile and not always well-defined. Nanobind has this information in a more structured form and could provide it.

3. Stubs serve two purposes, and stub generation should cater to both needs:

• To get autocomplete in VS Code and similar tools, which requires extracting function signatures and docstrings. I am mainly interested in this use case.

• Type checkers like MyPy. I haven't used them before and know very little (hence this issue to exchange experience). It seems to me that stubs only need to contain typed signatures but no docstrings. But nanobind often generates type annotations that MyPy isn't happy with, so some sort of postprocessing may be needed.

Here is what I have in mind, before having actually having done anything. There may be roadblocks I haven't considered.

1. The CMake build system gets a new command nanobind_add_stubs. This will register a command that is run at install time. Basically we need the whole package to be importable, and doing that in a non-installed build might be tricky.

nanobind_add_stubs(
  PATH ${CMAKE_INSTALL_PREFIX}
  PACKAGE nanobind_example
  DEPENDS nanobind_example_ext
)

When the user installs the extension to ${CMAKE_INSTALL_PREFIX}, this will run a Python file (shipped as part of the nanobind distribution) that imports the package and then generates nanobind_example/__init__.pyi.

Here, I am already getting confused because of unfamiliarity with stub generation. I've seen that packages sometimes contain multiple .pyi files. How does one decide where to put what? Can .pyi files import each other? What would be the best way to expose this in the nanobind_add_stubs() function?

2. I also wanted to modify nanobind's function class (nb_func) so that it exposes information in a more structured way, a bit like __signature__ from inspect.signature. But __signature__ is too limited because it (like Python) has no concept of overload chains.

Therefore, I am thinking of adding a function __nb_signature__ that returns list of pairs of strings [("signature", "docstring"), ...] that the stub generator can turn into something like this

from typing import overload

@overload
def func(x: int):
    """docstring 1"""

@overload
def func(x:str):
    """docstring 2"""

3. Some types signatures in nanobind aren't parseable in Python. There are a few things that I think could wrong:

• What if a C++ type hasn't been mapped yet when the extension is imported? In that case, the nanobind docstring includes the raw type (something like std::__1::vector<Item *>). In that case, the stubs could omit that overload entirely, put some generic placeholder (object?) or put the type name into a string. Thoughts?
• The representation of default arguments (via __repr__) might not make sense as a Python expression. This seems like an unsolvable problem because nanobind simply does not know the Python expression to re-create an object. One option would be to try to eval() the expression in the stub generator and omit it or replace it by some kind of placeholder if an issue is found. Not sure -- thoughts?
• Some nanobind type features don't have equivalents in typing.*. An example are the nd-array types annotations which are AFAIK too complex to be handled by anything currently existing. I'm thinking that it could be useful if the stub generator command nanobind_add_stubs(..) could be called with a user-provided Python file that implements some kind of post-process on the type signatures.

I'm curious about your thoughts on this! Thanks!

[jdumas]

Great to see some discussion on first-party support for stub generation in Nanobind.

I would like to add an important use-case for stub generation: documentation. As an example, we've generated this reference documentation for Lagrange, which uses Nanobind to create Python bindings (still very much WIP). I tried to use Sphinx Autodoc, but it doesn't support native C extensions (sphinx-doc/sphinx#7630). So instead I am using this Sphinx AutoAPI extension, which works by parsing stub files generated by Nanobind.

[wjakob]

@jdumas You might be interested in seeing the approach for an upcoming version of the Dr.Jit project that is based on nanobind: https://drjit.readthedocs.io/en/nanobind_v2/reference.html. I've so far managed without stubs -- this is the associated Sphinx file with plenty use of autofunction, autoclass, etc.: https://github.com/mitsuba-renderer/drjit/blob/nanobind_v2/docs/reference.rst?plain=1

[wjakob]

I started PR #421 to track progress towards landing this feature. The first commit (d71a990) adds a __nb_signature__ property to nanobind functions that exposes metadata about function overloads and default arguments in a more structured manner.

[njroussel]

1.

Here, I am already getting confused because of unfamiliarity with stub generation. I've seen that packages sometimes contain multiple .pyi files. How does one decide where to put what? Can .pyi files import each other?

Yes, somewhat.
Some autocomplete systems know how to deal with .pyi files importing other .pyi. As far as I can remember, this isn't properly documented by any PEP. In fact, if anything, the PEP 561 is pretty explicit about *.pyi files not requiring any "runtime" behavior. However, in practice, we do have such imports in https://github.com/mitsuba-renderer/drjit and they are picked up by Pyright (the underlying tool in VsCode and other IDEs).

What would be the best way to expose this in the nanobind_add_stubs() function?

I believe that splitting *.pyi files is only necessary if there are multiple submodules in the package. Finer-grain splitting is most likely only somewhat beneficial to the autocomplete system (higher parallelism when parsing). Exposing this control to the user would require them to provide a mapping from a stub file name to a list which contains all the symbol names they wish to have in said stub file. I really don't see an elegant way of handling this in CMake.

---

3

For the default arguments, there could be an extension to nb::arg_v that additionally takes a string which should be the Python expression to construct the default value. This is obviously extra work on the user, but as you said, there is no way of knowing the equivalent Python expression. I'm not against still have a default implementation that attempts to use eval(), but I think there should be an option for users to set it explicitly if they wish.

[wjakob]

@njroussel:
Regarding point 1, would it be sufficient to let the user specify a list of modules and submodules for which stubs should be generated, and then it makes separate file for each one?

After some thinking, the approach that I was planning regarding 3 is as follows (it somewhat mirrors what MyPy's stubgen does):

Simple default arguments that can be parsed out-of-the-box are included in the stub. More complicated things (elaborate types, multi-line strings) that might actually reduce readability of a function signature are abbreviated as "= ..." (where ... is the Python Ellipsis object).

The stub generator will by default accept types of: int, float, bool, NoneType, str (if not multi-line and shorter than a hardcoded limit) and nanobind enumerations. For more complex cases, I plan to have an extension mechanism. When calling the stub generator, the user can provide a custom Python file with a hook to turn other objects into Python expressions (or return None, which activates the default behavior). This seems less burdensome than having to define lots of custom nb::arg_v strings in a codebase (one place to catch them all, so to say).

Putting all together, it might look like this:

nanobind_add_stubs(
  # directory relative to which the stub generator will be called
  # (the modules should be importable from here)
  PATH ${CMAKE_INSTALL_PREFIX}
  # Sequence of modules and submodules, stubgen will separately process each one
  MODULES nanobind_example nanobind_example.submodule_1
  # Ensure that a given dependency is built and installed
  DEPENDS nanobind_example_ext
  # The user can provide a python file with hooks as part of their project
  # that modifies the behavior of stubgen
  HOOKS ${CMAKE_CURRENT_SOURCE_DIR}/stubgen_hooks.py
)

[wjakob]

cc also @skallweitNV & @laggykiller, who seem interested in stub generation

[skallweitNV]

This all sounds excellent. Some comments from my point of view:

I'd like to still be able to build stubs during development at build time. But people who want that can probably just call the stubgen manually from their build system. No need to have it in nanobind's cmake infrastructure.

Regarding default values. Handling them using hooks seems like a good idea, but you could still consider the nb::arg_v with default value string as an option. It might prove valuable in certain cases IMO.

And thanks for all your efforts @wjakob, nanobind is awesome!

[laggykiller]

Sounds great and looking forward for this to get implemented!

One thing to suggest, some user may add sphinx docstring to the nanobind code (e.g. myself, see this example). In those case, I think nanobind stub generator could try to parse sphinx docstring and use it as a hint for function signature and stub generation. This can provide a way for user to override auto generated stub if somehow the auto generated stub is incorrect.

Unfortunately, sphinx docstring do not have standardized way to denoate the default value of a function, so this can't be used to solve the problem of dealing with default arguments that do not make sense as a Python expression. Or, we could create our own standard of how default value of function should be denoted in sphinx docstring such that it could be recognized and used as hint during stub generation.

Another small thing to add is if representation of default argument do not make sense as a Python expression, we could denote the type of the argument as Optional as placeholder.

[rzhikharevich-wmt]

Another small thing to add is if representation of default argument do not make sense as a Python expression, we could denote the type of the argument as Optional as placeholder.

That's not really necessary as there is the ellipsis syntax for omitting default values from stubs. Besides, it would make type checkers validate code that explicitly passes an Optional value even when it's not applicable.

[wjakob]

Dear all,

I've made progress on this and have a working stub generation workflow that works well on a large-ish project (~300KB stub file with a mixture of classes, functions, properties, enums, and mixed bindings+pure Python code).

Could I ask you all to check it out and report any issue you encounter? The latest version is available on the stubgen branch.

There is also documentation on ReadTheDocs:

• Introduction: https://nanobind.readthedocs.io/en/stubgen/stubs.html
• CMake API: https://nanobind.readthedocs.io/en/stubgen/api_cmake.html#command:nanobind_add_stub

I followed @skallweitNV's suggestion and for now only have a CMake command that runs at build time. I plan to look into an install-time option later.

I also plan to make one major API and ABI-breaking change that will imply moving to nanobind version 2.0 (!)

The nb::raw_doc annotation completely overwrites a docstring of an overload chain. People usually use it to hand-craft a combined signature + docstring in cases where nanobind wasn't able to generate a usable/nice typed function signature.

m.def("f", &f, nb::raw_doc(R"(
my_func(a: complicated_type)

A docstring)"));

Interpreting this "raw" format is problematic when the stub generator needs to split up the signature into its constituent overloads. stubgen would the need to parse the docstring, and the whole point of this PR is to create an alternative approach for stub generation that doesn't involve parsing docstrings because it is brittle. Really the main reason for this raw format is to take control of the function signature.

Hence, my plan is to replace this with a new annotation named nb::signature. One can use it to specify multiple overloads, each with a custom signature (if needed) and docstring (if needed),

m.def("f", [](std::vector<int>) { }, nb::signature("(FantasticVector[int]) -> None"), "Docstring 1");
m.def("f", [](std::set<int>) { }, nb::signature("(FantasticSet[int]) -> None"), "Docstring 2");

stubgen will simply copy-paste the signature into a @typing.overload chain. In this example, it will generate

from typing import overload

@overload
def f(FantasticVector[int]) -> None:
    """Docstring 1"""

@overload
def f(FantasticSet[int]) -> None:
    """Docstring 2"""

del overload

[laggykiller]

Nice work! Some minor problems while trying out using my project: https://github.com/laggykiller/apngasm-python/tree/stubgen-official

---

Problem 1

I tried to added the following lines in CMakeLists.txt:

nanobind_add_stub(
    apngasm_python_stub
    MODULE apngasm_python._apngasm_python
    OUTPUT src-python/apngasm_python/_apngasm_python.pyi
    MARKER_FILE src-python/apngasm_python/py.typed
    VERBOSE
)

This does not work if apngasm_python is not installed while wheel is being built as apngasm_python cannot be imported for stub generation. As a workaround, I have to generate stub file manually after building and installing:

python -m nanobind.stubgen -m apngasm_python._apngasm_python -o src-python/apngasm_python/_apngasm_python.pyi -M src-python/apngasm_python/py.typed

This is not quite ergonomic. If I want to build a wheel with stub, I have to build and install wheel without stub first, run the command to generate stub, then build the wheel with stub!

---

Problem 2

Reference to classes within the same file is not correct. Consider this pyi file

class APNGAsm:
    @overload
    def __init__(self, frames: Sequence[apngasm_python._apngasm_python.APNGFrame]) -> None:
        ...

class APNGFrame:
    ...

Pylance would get lost:

It should be:

class APNGAsm:
    @overload
    def __init__(self, frames: Sequence[APNGFrame]) -> None:
        ...

class APNGFrame:
    ...

Pylance is happy with this:

Alternatively:

from . import _apngasm_python

class APNGAsm:
    @overload
    def __init__(self, frames: Sequence[_apngasm_python.APNGFrame]) -> None:
        ...

class APNGFrame:
    ...

---

Problem 3

It is not necessary to delete imports in pyi, it triggers warning in pylance:

[laggykiller]

P.S.

Consider this function to be binded

size_t addFrame(rgb *pixels, unsigned int width, unsigned int height, rgb *trns_color = NULL, unsigned delayNum = DEFAULT_FRAME_NUMERATOR, unsigned delayDen = DEFAULT_FRAME_DENOMINATOR);

Note that trns_color has type rgb* and default value is NULL.

This is the binding code

        .def("add_frame_from_rgb", nb::overload_cast<apngasm::rgb *, unsigned int, unsigned int, apngasm::rgb *, unsigned, unsigned>(&apngasm::APNGAsm::addFrame),
        "pixels_rgb"_a, "width"_a, "height"_a, "trns_color"_a = NULL, "delay_num"_a = apngasm::DEFAULT_FRAME_NUMERATOR, "delay_den"_a = apngasm::DEFAULT_FRAME_DENOMINATOR)

This is the signature generated

 def add_frame_from_rgb(self, pixels_rgb: apngasm_python._apngasm_python.rgb, width: int, height: int, trns_color: apngasm_python._apngasm_python.rgb = 0, delay_num: int = 100, delay_den: int = 1000) -> int:

I would expect trns_color signature to be trns_color: Optional[rgb] = None

[wjakob]

@laggykiller Thanks for giving it a try. The challenge with generating stubs at build vs install time is mentioned above and was something I still wanted to look into.

In your example, wouldn't it be better to generate a stub for the actual module containing your types? (i.e., apngasm_python._apngasm_python). Then you could add a manual stub file __init__.py that imports the types.

Alternatively let me teach you about a dirty hack that I use in my own projects.

NB_MODULE(my_module_ext, m_) {
   (void) m_; /* unused */
   nb::module m = nb::module_::import_("my_module");
   // register functions and classes here
}

In other words, importing the sub-extension causes it to add its types to the parent module. That way, all of the type metadata (__module__) causes the created functions and types to be in the expected place as seen by consumers of the API.

An alternative option for the stubs would be to register a rewriting rule with the stub generator that rewrites all occurrences of apngasm_python._apngasm_python to apngasm_python. I am thinking about adding the capability for such custom hooks but have not done so yet.

For your second message, does it work if you annotate the argument as "trns_color"_a.none() = NULL?

The del declarations at the end are caused by my inexperience with stubs in general. At a high level, it does make sense that they are there, no? We don't want Optional, Annotated, etc., to be listed as official exports of the stub. How does one do this then? Import them as _Optional with an underscore?

Edit: never mind about the last point. I found out that the typing documentation explains a clear set of rules when an import is considered to be part of the stub exports.

[laggykiller]

wouldn't it be better to generate a stub for the actual module containing your types?

I don't quite get what you are trying to say, but I think my command is doing what you said...? I specified -m apngasm_python._apngasm_python when calling python -m nanobind.stubgen command.

---

Alternatively let me teach you about a dirty hack that I use in my own projects.

I replaced the following lines in my binding code:

NB_MODULE(MODULE_NAME, m) {
    m.doc() = "A nanobind API for apngasm, a tool/library for APNG assembly/disassembly";

Replaced as:

NB_MODULE(MODULE_NAME, m_) {
    (void) m_; /* unused */
    nb::module m = nb::module_::import_("apngasm_python");
    m.doc() = "A nanobind API for apngasm, a tool/library for APNG assembly/disassembly";

This causes error during compilation:

error: ‘module’ is not a member of ‘nb’; did you mean ‘module_’?
   37 |     nb::module m = nb::module_::import_("apngasm_python");
      |         ^~~~~~
      |         module_

So I replaced nb::module m as nb::module_ m. Compilation was 'successful', but now apngasm_python._apngasm_python is empty!

---

does it work if you annotate the argument as "trns_color"_a.none() = NULL?

The signature generated is trns_color: Optional[apngasm_python._apngasm_python.rgb] = 0. Optional is added correctly, but Pylance is unhappy that the default value is 0 instead of None.

---

At a high level, it does make sense that they are there, no? We don't want Optional, Annotated, etc., to be listed as official exports of the stub.

I think it is fine to not del imports, see a random example from Microsoft: https://github.com/microsoft/python-type-stubs/blob/e328827329a072baac06e287bd6146ff80796e5e/stubs/skimage/transform/hough_transform.pyi

---

Btw, would it be better that instead of:

from typing import Annotated
from typing import Optional
from typing import Sequence
from typing import overload

We import by...

from typing import Annotated, Optional, Sequence, overload