#6399·pyo3

Type stubs customization

Author: TptCreated Sep 9, 2026Updated Sep 11, 2026

This is a bold proposal triggered by discussions on #5877. It's more than likely to be missing some very important things and not be usable in practice. The main goal here is to explore a bit more the problem and solution spaces.

Context

It seems to me user-provided type stubs are useful in 3 cases:

  1. In case of custom type annotations for function input parameter or output value, to get the corrects import. For example, if the custom annotation is Sequence[int] we need from collections.abc import Sequence import to make it correct.
  2. FromPyObject and IntoPyObject implementations (both manual or derived) and custom type annotations might want to leverage Protocols (and their specialized variant TypedDict). For example if we accept any object with a field name foo which value must be an int, we need a Protocol to express that. A good example of that is "capsule interfaces" like the Arrow one.
  3. In case of custom addition to modules with #[pymodule_init] functions. Content can be anything (classes, constants, reexports…).

We might solve the 3 of them by allowing custom type stubs at the module level. However, positioning these custom stubs alongside the existing ones generates a "chicken and egg" problem: these custom stubs might require some types declared using PyO3 in the same module and, hence, in the auto generated stubs question but might want to expose imports for custom type annotations.

For example, in

rust
#[pymodule]
mod foo {
    #[pyfunction(signature = (a: "time") -> "time")]
    fn cp_time(a: Bound<'_, PyAny>) -> Bound<'_, PyAny> { a }

    #[pymodule_init]
    fn module(m: &Bound<'_, PyModule>) -> PyResult<()> {
         m.set_attr("cp_time_alias", m.get_attr("cp_time")?)?;
    }

We might want the custom stubs to be:

python
from datetime import time
cp_time_alias = cp_time

However the first line must be before the def cp_time(a: time) -> time: ... declaration and the second one after.

Also, these module-level type stubs don’t solve nicely the use case of a Protocol useful for a FromPyObject implementation done in a library crate and used in other crates because all these other crates would need to include the protocol in their custom type stubs, leading to subpar UX. So, it might be better to solve the 3 use cases separately.

Proposal

1. custom annotations

We introduce a Rust DSL for annotations in #[pyo3(signature = ...)] allowing to write #[pyo3(signature = () -> collections.abc.Sequence[int])] instead of #[pyo3(signature = () -> "collections.abc.Sequence[int]")] (now without quotes!). We auto-generate imports by considering all but the last element of the paths to be the module. In our example collections.abc is the module and Sequence the imported type. This should cover all use cases except nested classes. For nested classes we might introduce an abusing notation asking people to write something (foo.bar).Class.SubClass with the parentheses highlighting the module path (here foo.bar).

2. protocols

This is the most painful part. We need to allow expressing protocols with an arbitrary number of methods, getters, setters, dictionary magic methods with @override for dicts (__getitem__, __set_item__...). Parsing Python syntax in Rust macros is a pain so a bespoke DSL might be better (importing type stubs snippet verbatim will be painful because we want to ensure things like no name conflicts if two FromPyObject implementations want to use the same protocol name...). A syntax idea, building on top of the existing #[pyo3(signature = )] syntax:

{
    method(foo: int) -> int;
    @getter
    foo(self) -> datetime.time;
    __getitem__(name: typing.Literal["name"]) -> str;
}

This DSL would be usable from two places:

  • using a macro type_hint_protocol!("MyProtocolName", { ... }) where ... is the protocol definition. The macro returns a PyStaticExpr to be used in explicit FromPyObject and IntoPyObject implementations. The name "MyProtocolName" is just a hint for the generator, it can be changed into e.g. _MyProtocolName2 in case of conflicts.
  • inside of #[pyo3(signature = )] like #[pyo3(signature = (a: { method(foo: int) -> int }) -> int)]. The name of the protocol is automatically generated in this case (something like _FunctionNameArgNamePrococol).

3. Extra module stubs

We have solved the import and protocol problems for types declared by PyO3 hence we can just put the user provided type stubs at the end of the files. Because the PyO3 type stubs output is deterministic, we might just let the user handle possible name conflicts between the PyO3 generated stubs and their custom stubs. We might just add a extra_stubs parameter to #[pymodule] that is either a path to a .pyi file to be included or some inline stubs.