Type stubs customization
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:
- In case of custom type annotations for function input parameter or output value, to get the corrects import.
For example, if the custom annotation isSequence[int]we needfrom collections.abc import Sequenceimport to make it correct. FromPyObjectandIntoPyObjectimplementations (both manual orderived) and custom type annotations might want to leverageProtocols (and their specialized variantTypedDict).
For example if we accept any object with a field namefoowhich value must be anint, we need aProtocolto express that.
A good example of that is "capsule interfaces" like the Arrow one.- 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 aPyStaticExprto be used in explicitFromPyObjectandIntoPyObjectimplementations.
The name"MyProtocolName"is just a hint for the generator, it can be changed into e.g._MyProtocolName2in 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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