PyO3 / PyO3/pyo3

Type stubs customization

Open
#6,399 4 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Rust
Stars
16.2k
Forks
1k
Avg merge
2d 6h
Merged PRs (30d)
66

Description

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

#[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:

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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the discussion referenced in #5877 and the existing #[pyo3(signature = ...)] and type-stub generation behavior. Compare the three proposed areas—custom annotations, protocols, and extra module stubs—before narrowing the scope. Done would require an agreed design and implementation criteria, which this proposal does not yet define.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, rust
Domain
tooling
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.