Python Interface (.pyi) generation and runtime inspection

[CLOVIS-AI]

Hi!

This issue is going to be a summary of my prototypes to generate Python Interface files ('stub files', '.pyi files') automatically. The prototypes are available as #2379 and #2447.

• [Prototype] Python stubs generation #2379 aimed to generate Python stub files entirely at compile-time. This is not possible because proc-macros run before type inference and trait analysis, so proc-macros cannot know if a type implements a trait or not.
• Prototype: Runtime inspection and Pyi generation #2447 aims to generate information at compile-time that represents the structure of #[pyclass] structs (which methods exist, what arguments do they accept) to be read at run-time by the stub generator.

I'm presenting the results here to get feedback on the current approach. I'm thinking of extracting parts of the prototypes as standalone features and PRs.

Progress

Accessing type information at runtime:

• Runtime type information for objects crossing the Rust–Python boundary #2490
• Generate precise type information as part of #[pyclass]

Accessing structural information at runtime:

• Declare the inspection API
• Generate the inspection API structures as part of #[pyclass] and #[pymethods]
• Collect the inspection data per module

Python interface generation:

• Generate the Python interface files
• Document how to generate PYI files in end user's projects

Summary

The final goal is to provide a way for developers who use PyO3 to automatically generate Python Interface files (.pyi) with type information and documentation, to enable Rust extensions to 'feel' like regular Python code for end users via proper integration in various tools (MyPy, IDEs, Python documentation generators).

I have identified the following steps to achieve this goal. Ideally, each step will become its own PR as a standalone feature.

1. provide a way to extract the full Python type information from any object passed to/retrieved from Python (e.g. List[Union[str]], not just PyList).
2. provide an API to describe Python objects at run-time (list of classes, list of methods for these classes, list of arguments of each method, etc).
3. improve the macros so they generate at compile-time the various inspection data structures (the API from 2.)
4. write a run-time pyi generator based on the inspection API

1 and 2 are independent, 3 and 4 are independent.

Full type information

The goal of this task is to provide a simple way to access the string representation of the Python type of any object exposed to Python. This string representation should follow the exact format of normal Python type hints.

First, a structure representing the various types is created (simplified version below, prototype here):

struct TypeInfo {
    Any,
    None,
    Optional(Box<TypeInfo>),
    ...
    Builtin(&str),
    Class {
        module: Option<&str>,
        name: &str,
    }
}

impl Display for TypeInfo {
    // Convert to a proper String
}

PyO3 already has traits that represent conversion to/from Python: IntoPy and FromPyObject. These traits can be enhanced to return the type information. The Python convention is that all untyped values should be considered as Any, so the methods can be added with Any as a default to avoid breaking changes (simplified version below, prototype here):

pub trait IntoPy<T> {
    // current API omitted

    fn type_output() -> TypeInfo {
        TypeInfo::Any
    }
}

pub trait FromPyObject {
    // current API omitted

    fn type_input() -> TypeInfo {
        TypeInfo::Any
    }
}

The rationale for adding two different methods is:

• Some structs implement one trait but not the other (e.g. enums which use derive(FromPyObject)), so adding the method to only one of the trait would not work in all cases,
• Creating a new trait with a single method would be inconvenient for PyO3 users in general, as it would mean implementing one more trait for each Python-exposed object
• Both methods have a sensible default, and are both trivial to implement so I don't believe there are any downsides,
• Some Python classes should have a different type when appearing as a function input and output, for example Mapping<K, V> as input and Dict<K, V> as output. Using two different methods supports this use case out-of-the-box.

After this is implemented for built-in types (prototype here), using them becomes as easy as format!("The type of this value is {}", usize::type_input()) which gives "The type of this value is int".

Inspection API

This section consists of creating an API to represent Python objects.

The main entry point for users would be the InspectClass trait (simplified, prototype here):

pub trait InspectClass {
    fn inspect() -> ClassInfo;
}

A similar trait would be created for modules, so it becomes possible to access the list of classes in a module.
This requires creating a structure for each Python language element (ModuleInfo, ClassInfo, FieldInfo, ArgumentInfo…, prototype here).

At this point, using this API would require instantiating all structures by hand.

Compile-time generation

Proc-macros can statically generate all information needed to automatically implement the inspection API: structural information (fields, etc) are already known, and type information can simply be delegated to the IntoPy and FromPyObject traits, since all parameters and return values must implement at least one of them.

Various prototypes:

• 38f0a59: extract classes
• 56b85cf: extract the list of functions
• 8125521: extract a function's kind (function, class method, static method…)
• 4070ad4: extract the function's return type,
• 53f2e94: extract attributes annotated with #[pyo3(get, set)],
• 003d275: extract argument names and type

This is done via two new traits, InspectStruct, InspectImpl which respectively contain the information captured from #[pyclass] and #[pymethods]. Due to this, this prototype is not compatible with multiple-pymethods. I do not know whether it is possible to make it compatible in the future.

Python Interface generator

Finally, a small runtime routine can be provided to generate the .pyi file from the compile-time extracted information (prototype here).

Thanks to the previous steps, it is possible to retrieve all information necessary to create a complete typed interface file with no further annotations from a user of the PyO3 library. I think that's pretty much the perfect scenario for this feature, and although it seemed daunting at first, I don't think it's so far fetched now 😄

The current state of the prototype is described here: #2447 (comment).

[davidhewitt]

@CLOVIS-AI just a ping to say I haven't forgotten about this; have been ill / busy and was just about cleared through the backlog enough to read this when #2481 came up. I think I need to push through some security releases first, after which my plan was to finish #2302 and then loop back here with us ready to support a syntax for annotations. Sorry for the delay.

[CLOVIS-AI]

@davidhewitt don't worry about the full review for now, it's just a prototype. If you have time, please just read this issue and give me your general feedback on the idea. If it seems good with you, I'll be able to start writing a real PR for at least a part of it and we can do a full review then 👍

[davidhewitt]

Hey, so I finally found a moment to sit down and think about this. Thank you for working on this and for having patience with me.

This looks great, I think this is definitely the right way to go. In particular splitting into two traits for the input/output I think is correct.

Some thoughts and questions:

• I've wanted something like FromPyObject::type_input for a long time, I think it can be used to implement improved error messages. In particular in PyDowncastError we currently store the type name, e.g. PyString, but with type_input we could do something better (as long as the input is not Any, I guess).
• I think it would be nice to feature-gate the inspection API and macro codegen for it - presumably it would be used with a debug build to emit type stubs, and then it wouldn't be needed in the final release build.
• To be valid .pyi files they often need to import type dependencies. Do you have a vision how we might be able to make this work?
• Imagine users create custom MyDict[K,V] class which would be a generic type (in Python's eyes, the Rust code would potentially just use PyAny). Can we support it with this proposal?
• For cases like the above, if we can't support them, we might need mechanisms to load external type stub fragments to combine into the final artifact.

Overall, yes, I'm happy to proceed with this - as a first step I'd suggest we get type_input and type_output merged. We could already use them to improve error messages and also maybe add function signatures in their docstrings. That would buy us some time to figure out the introspection API, which I think will have some complexity.

[CLOVIS-AI]

I was thinking of feature-gating the macro generation but not the inspection API itself (so you would be able to construct inspection APIs yourself in all cases, but would need to enable the automatic implementation). I assume that the API itself will not have any significant effect on compile-time, since it's just normal structs. What do you think?

The TypeInfo struct has a module method that returns the name of the module a class is from. When generating a .pyi file, a recursive visit of all types appearing in the file would yield the list of modules to import. This may be too simplistic (and will probably fail if two classes in different modules have the same name), but it should be good enough for most users (for example, in my case, we have a single module anyway).

About custom generics: the approach in this documentation couldn't, but the one in #2490 can (however, it must be through user input, I don't see a way in which the macros could guess that).

Combining external information with the generated ones will be trivial: because the macros will generate an implementation of the inspection API, and the .pyi generation will take an implementation as parameter, users can simply edit the generated implementation before passing it to the .pyi generation.

[CLOVIS-AI]

It seems like #2490 will be merged soon. I won't have a lot of time on my hands in the close future, so if someone else wants to help in the meantime, the next big question is the way to represent the program (Python classes, Python methods, Python modules) as Rust structures.

My prototypes are close to solving the problem, except that I'm not a fan of how they deal with modules. The structures themselves seem fine, but the way to convert a #[pymodule] function into them is unclear.

[Tpt]

Not sure if it relevant here:

I have written a small python script to generate type stubs from pyo3 libraries with doc strings including type annotations (using the :type and :rtype: format).
It works on the already built libraries using the Python introspection feature.
To use it run python generate_stubs.py MY_PACKAGE MY_FILE.pyi
Here is an example of generated stubs.

[CLOVIS-AI]

@Tpt That's great! However if I understand correctly you still have to declare the type twice (first as a Rust type, then as a Python type in the documentation), which is error-prone, and what this issue tries to avoid. I agree that it's already a great step up from the current situation of writing the .pyi entirely manually.