pyo3

pyo3 copied to clipboard

This is a first step to introspect Python modules built by PyO3

A missing piece in the story listed in #2454 is how tools like Maturin move the introspection information generated how by PyO3 into to type stubs files included in the built wheels.

I see three approaches for it:

1. Make pyo3-macros generate a file with the stubs after having processed all macros of the crate. This has the advantage of being self-contained in the crate but falls short in cases like python classes declared in a crate but exposed in an other crate: there is no guarantees that proc macros of a crate and its dependencies are compiled in the same process and that proc macros will still be able to write files in the future (like with proposal to run them a WASM sandbox).
2. Make pyo3-macros exports a C function that returns the stubs. The built libraries can then be loaded by Maturin that would call the function and write the stubs to a file. I wrote a quick experiment with a patch to PyO3 that generate a public C ABI __pyo3_stubs_MODULE_NAME function However, for the build system to execute it, a compatible Python interpreter must be present to link with and a compatible CPU or VM to run it, making generation when doing cross-compilation very hard. I guess it's what #2454 was heading toward.
3. Add the introspection data in a custom section of the built cdylib. It requires the introspection data to be completely static. However, we can easily build a library that extract and parses these sections and generate the stub file. Build tool like Maturin would just need to call this library to get the stubs. Cross compilation support is easy as soon as a parser for the built cdylib format exists. However, this might make the generated binaries bigger because of the extra data. This is what the MR experiments with.

Architecture:

• Each PyO3 element (pymodule, pyclass, pyfunction...) generates a segment in the built binary pyo3_data0 section that contains a JSON "chunk" with the introspection data. Code in pyo3-macros-backend/src/introspection.rs. I had to do some bad hack to generate the segments properly via Rust static elements.
• JSON chunks can refer to other chunks via global ids. These ids are stored in the PYO3_INTROSPECTION_ID constants, allowing the code building the JSON chunk to get the global id of eg. a class C via C::PYO3_INTROSPECTION_ID. This allows chunks to refer to other chunks easily (eg. to list module elements). A bad hack is used to generate the ids (TypeId::of would have been a nicer approach but is not const on Rust stable yet).
• The 0 at the end of pyo3_data0 is a version number, allowing breaking changes in the introspection format.
• The pyo3-introspection crate parses the binary using goblin (library also used by Maturin), fetches all the pyo3_data0 segments (only implemented for macOS Match-O in this experiment), and builds nice data structures.
• Not done yet: the pyo3-introspection crate would implement a to_stubs function converting the data structures to Type stubs.
• The pyo3-introspection has an integration tests doing introspection of the pyo3-pytests library.

Current limitations:

• WASM is not supported yet, only PE (Windows), Mach-0 (macOS) and ELF (Linux and some other *nix).
• the only introspection data created is the list of function, classes and submodules inside of a module created using the new declarative module syntax
• this approach requires to convert FromPyObject::type_input into an associated constant or a const function and similarly for IntoPy::type_output. This is mandatory in order to allow to make use of them in the static values added to the generated binary.

Mar 22 '24 09:03 Tpt

Thanks for moving this forward! The idea of using custom data sections is new to me. I see the upsides of it, though I am slightly worried by the extra complexity of needing to worry about linker details in yet another way.

1. Make pyo3-macros generate a file with the stubs after having processed all macros of the crate.

I agree that having the macros generate file(s) is unlikely to be the right solution 👍

2. However, for the build system to execute it, a compatible Python interpreter must be present to link with and a compatible CPU or VM to run it, making generation when doing cross-compilation very hard. I guess it's what Python Interface (.pyi) generation and runtime inspection #2454 was heading toward.

For the library which converts the metadata to .pyi files, I think rather than expecting the build system like maturin to execute it, it would work better if we asked users to keep their .pyi file committed in the repository and they used the library to generate / update the file as part of their test suite. I think in that case then local development should be enough and cross-compiling can be ignored from the picture.

In fact, I wonder if for option (3) explored here then the using a test to generate and update the .pyi file is still better than to run it during maturin build. Here's why:

• I think it's easier for users to understand what their .pyi file contains if it's committed in the repository rather than generated by maturin as part of the packaging process.
• With a test, users are hopefully reminded in CI if they forgot to regenerate their .pyi files.
• No need for us to include this code / information from the binary if it's only needed during development. (Probably we use a cargo feature which users only enable in their dev-dependencies?)

If we agree that using a test to update stubs is a good solution, then I think the choice between runtime code like (2) and data segments like (3) is probably just influenced by whatever is easier for us to implement. We might even be able to swap back and forth between these two options as an implementation detail as we learn.

Mar 23 '24 22:03 davidhewitt

If we agree that using a test to update stubs is a good solution, then I think the choice between runtime code like (2) and data segments like (3) is probably just influenced by whatever is easier for us to implement. We might even be able to swap back and forth between these two options as an implementation detail as we learn.

Yes! To have played a bit with both, runtime code like (2) is way easier (the difference between this MR and #2454 is quite significant).

If I try to summarize the pros of each approach:

Approach 3: add introspection to the cdylib and let maturin write the stubs on build:

• no extra user code/configuration: they upgrade maturin and pyo3 and it just works instead of having to setup a test, make extension-module not a default...
• support of Rust features: if a feature is enabled, extra stubs might get generated and if it is disabled stubs won't get generated. This seems impossible to get with approach 4
• easier Python version and platform customization: stubs are generated per Python version and platform

Approach 4: use a test to write/update stubs:

• hopefully simpler code in PyO3 (not data segment generation and retrieval from binary...)
• easier debugging and change inspection: stubs are commited as part of repository content
• Python version and platform specific elements are explicit in the stubs. But it might be painful to implement because it might be a bit of cat-and-mouse game with Rust #[cfg] annotation.
• smaller generated cdylib (no introspection data in them)

Mar 25 '24 09:03 Tpt

You make a very good point that automated tests to update a .pyi file struggle with cfg declarations. I suppose one way around that would be to make it possible to have multiple .pyi files committed for different Python / OS combinations and select the correct one in packaging.

One thing that I expect is that .pyi files might want to include extra user customisations beyond what the autogenerated stubs contain. For example, I expect generic code may be tricky to define on the Rust side and may benefit from some handwritten customisation of stubs. (e.g. making a user-defined dict[K, V] class.) I don't have a good answer to how we can solve that.

Overall I don't have a good sense of whether option 3 or 4 is better. In an ideal world we might offer both options. Which one do you think would meet your needs better at present? Maybe we start by implementing that and we learn a lot by doing so!

Mar 26 '24 23:03 davidhewitt

One thing that I expect is that .pyi files might want to include extra user customisations beyond what the autogenerated stubs contain. For example, I expect generic code may be tricky to define on the Rust side and may benefit from some handwritten customisation of stubs. (e.g. making a user-defined dict[K, V] class.) I don't have a good answer to how we can solve that.

I would love to avoid people to handwrite customization in stubs because it makes automatically updating the stubs when Rust code is changed very hard. Imho automatically updating stubs to reflect changes in the Rust code is the main value proposition of auto-generating stubs in the first place.

An idea: Add entry points in PyO3 macros to extend the stubs. For example (rought idea, not sure about the actual details):

#[pymodule]
#[py(extra_stub_content = "
K = VarType('T')
V = VarType('V')
")]
mod module {
    #[pyclass(stub_parent = [collections::abc::Mapping::<K,V>])]
    struct CustomDict;
    
    #[pymethods]
    impl CustomDict {
       #[pyo3(signature = (key: K), return_type = V | None)]
       fn get(key: &Bound<'_, PyAny>) -> Option<Bound<'_, PyAny>> {
       }
    }
}

would generate

K = VarType('T')
K = VarType('V')

class CustomDict(collections.abc.Mapping[K,V]):
     def get(key: K) -> V | None: ...

This way stubs would stay auto generated but can be improved by the author.

A possible way to mix options 3 and 4:

• When building a wheel, Maturin uses by default the present stubs file or, if not present, falls back to automated introspection (option 3)
• Maturin provides a generate-stubs command that takes for input arguments like --target and --features and generates stubs file, allowing users to pick option 4 by commiting the command result to git. Users are also able to use this command to check that commited to git are up to date (run the command again and check if anything changed) or run some validation on the stubs in the CI using eg. mypy.stubtests even if the stubs are not commited to git.

Mar 27 '24 09:03 Tpt

Agreed that having the proc macros be able to collect all the necessary information would be nice. I think only time will tell whether they can meet all user needs!

I'm slightly wary of coupling to maturin for all stub generation, because some projects use setuptools-rust or their own build options for good reasons. I think that offering maturin generate-stubs as an alternative development command would be a good way to avoid the problem for projects that don't want to do their packaging with maturin.

cc @messense do you see any concerns with adding this to maturin?

So what's next steps here? Do you want me to start reviewing this code, or will you push more first?

Regarding the data sections, I happened to hear yesterday that UniFFI's proc macros can do something similar about shipping definitions in the shared library, so it might be interesting to look at / ask them how that was implemented.

Mar 27 '24 11:03 davidhewitt

do you see any concerns with adding this to maturin?

No concern, I think a generate-stubs command will be very useful for users wanting to commit pyi files in git. We can also add a --check option to fail the command when existing pyi file is outdated.

Mar 27 '24 11:03 messense

Thank you!

Agreed that having the proc macros be able to collect all the necessary information would be nice. I think only time will tell whether they can meet all user needs!

Yes! My hope is to cover as many as possible.

I'm slightly wary of coupling to maturin for all stub generation, because some projects use setuptools-rust or their own build options for good reasons. I think that offering maturin generate-stubs as an alternative development command would be a good way to avoid the problem for projects that don't want to do their packaging with maturin.

Additionaly to maturin generate-stubs, I would go one step further and suggest we also follow the approach started in this MR, ie. build a pyo3-introspection stand-alone crate that contains all the code to generate the actual stubs from the binary annotations. This way other build systems would be able to ship integrated automated stub generation if they want. This would also enable other tools like doing breaking change checking based on the introspection data.

Regarding the data sections, I happened to hear yesterday that UniFFI's proc macros can do something similar about shipping definitions in the shared library, so it might be interesting to look at / ask them how that was implemented.

Thank you! I'm going to have a look at it.

So what's next steps here? Do you want me to start reviewing this code, or will you push more first?

I think the current draft already shows the relevant direction, a very high level code review to check if it's going in the good direction would be welcome. Maybe wait for me to have a look at UniFFI, I might change a bit this MR if I find interesting things there. Thank you!

Mar 27 '24 17:03 Tpt

This is very exciting, looking forward to being able to generate type stubs! Currently we have this lengthy and hard to maintain Python script for doing so, which we have to update by hand: https://github.com/Chia-Network/chia_rs/blob/main/wheel/generate_type_stubs.py

This would be a major improvement. Happy to help out however I can (testing, implementation, whatever) as time allows, to hopefully get this out the door 😄

Mar 27 '24 18:03 Rigidity

@Rigidity Thank you! I plan to work on this MRs to get the basics done. Then there will be a lot of features to incrementally add on it (support for all PyO3 features...) so help will coding and testing will be much welcome!

Mar 27 '24 20:03 Tpt

Sorry for the very long reaction delay (a lot of priorities + vacations).

Regarding the data sections, I happened to hear yesterday that UniFFI's proc macros can do something similar about shipping definitions in the shared library, so it might be interesting to look at / ask them how that was implemented.

I had a look at uniffi, they basically use the same approach as us: embedding the metadata in the binary and then parsing the binary using the same goblin library. The major different I see is that my prototype is putting the metadata into a custom section named set using link_section whereas uniffi is using regular const without custom link_section. I would tend to think custom link_section is a bit clearer when analyzing the binary but I am not sure how much it changes things.

@davidhewitt If you have time, may you have a quick look at the MR to see if the global design goes into a good direction? If yes, I will fix a lot of shortcut I took and get the MR ready for review.

May 10 '24 15:05 Tpt

I would tend to think custom link_section is a bit clearer when analyzing the binary but I am not sure how much it changes things.

I think so too; I also imagine we might want to strip the custom link section after the stubs have been extracted, it feels like it's probably easier to do that by having a distinct section.

Jun 08 '24 07:06 davidhewitt

I'm definitely convinced by technical direction here for the generation process; the main risk I still see is how to give users the full power to customise the stubs with generic args etc.

I think we can learn by that piece-by-piece as we proceed.

Thank you! I agree on the risk. My guess is that we will introduce a set of macro arguments for that, but getting them right won't be easy.

Jun 10 '24 06:06 Tpt

I have implemented introspection for ELF and PE binaries covering Windows, macOS, Linux and hopefully various *nix.

However, I encounter an issue with WASM: "statics with a custom #[link_section] must be a simple list of bytes on the wasm target with no extra levels of indirection such as references". This leaves me with two choices:

• Do not use a custom #[link_section] but rely on specific symbol names like uniffi.
• Keep a custom #[link_section] but instead of having const _: &str to define the introspection data, use const _: [u8; i]. This means that the macros will have to compute the i, making the generating the introspection data a bit more cumbersome.

I am going to investigate the second option, if I can hide the complexity behind a fairly good interface, it might be good.

Edit: the second option is not a good idea: for example, it prevents to add to the introspection data type name that would be provided by e.g. FromPyObject::input_type

Jun 11 '24 08:06 Tpt

I have implemented introspection for ELF and PE binaries covering Windows, macOS, Linux and hopefully various *nix.

However, I encounter an issue with WASM

Isn't type stub generation only something that would have to happen on a development host machine? Maybe cross platform support would be fine to be somewhat limited as long as it's not done in regular builds

Jun 11 '24 10:06 Rigidity

Isn't type stub generation only something that would have to happen on a development host machine? Maybe cross platform support would be fine to be somewhat limited as long as it's not done in regular builds

The challenge with this approach is that it is very hard in Rust to know what are all the elements (class, function...) that might be exposed to any platform because code disabled for a plateform is just not compiled and its macros are not evaluated. Hence, if a library has some WASM-specific code, it is very hard to build introspection data for it while doing a Linux build. Hence, I went into the direction where we would generate stubs for each plateform at compile time. Distribution-side, it would mean that the source .tar.gz achive would not contain stubs but built .wheel would.

Jun 11 '24 11:06 Tpt

@Tpt Do you need help with implementing/testing this? Or you pretty much done and waiting for review? I'm willing to help push this MR, as it is much needed feature for my project. Thanks for your awesome work!

Aug 13 '24 09:08 yogevm15

@yogevm15 A basic MVP (generating stubs with just the list of classes and functions) is done. I would love to get review on it before adding more features to avoid a huge MR. I already got some nearly finished code on my laptop for function signatures but a lot of work still needs to be done for class instrospection. Help will be very welcome in this area when this MR is merged

Aug 14 '24 09:08 Tpt

Sorry for the huge delay from me on this one. I am generally beginning to catch up with the backlog from paternity leave; I will try to review this ASAP in the coming days so that we can unblock and move forward.

Aug 16 '24 21:08 davidhewitt

@davidhewitt Congratulations! Take your time, babies grow up so fast. Thank you!

Aug 17 '24 19:08 Tpt

@Tpt I was trying to test your feature, but I couldn't make it work with debug builds:

thread 'main' panicked at /pyo3/pyo3-introspection/src/introspection.rs:212:38:
range end index 72057675624348367 out of range for slice of length 988960

I am on macOS.

use pyo3::prelude::*;

#[pyclass]
struct DummyClass {}

#[pymethods]
impl DummyClass {
    #[staticmethod]
    fn get_42() -> PyResult<usize> {
        Ok(42)
    }
}

#[pymodule]
pub mod pyo3_pure {
    #[pymodule_export]
    use super::DummyClass;
}

use std::env;
use std::path::PathBuf;
use pyo3_introspection::{introspect_cdylib, module_stub_files};

fn main() {
    let binary = PathBuf::from(env::args().nth(1).unwrap());

    let module = introspect_cdylib(binary, "pyo3_pure").unwrap();
    let actual_stubs = module_stub_files(&module);

    dbg!(actual_stubs);
}

Aug 19 '24 08:08 Bben01

@Bben01 Thank you! Indeed, extraction was bugging on Match-O in dev. Fixed in 77e1b4d63b0923a23c0b98c77a7796a164e1eaa8

Aug 19 '24 10:08 Tpt