pyo3 icon indicating copy to clipboard operation
pyo3 copied to clipboard

add `ModuleState`, core lifecycle callbacks for `#[pymodule]`s

Open bazaah opened this issue 2 months ago • 5 comments

This patch set adds the core internal representation for a PyO3 module's per-module state.

Some of this work is taken from #4162 (credit: @Aequitosh), stripped down.

Within we...

  1. Added the core repr(C) type that CPython will allocate (space for)
  2. Add the requisite hooks in PyModuleDef for initialization (m_slots Py_mod_exec slot) and freeing (m_free)
  3. Add a TypeMap based store for user state types and expose a limited API to this state via the PyModuleMethods trait
  4. Test module initialization works as expected

Left for future work:

  1. Performance optimizations (we currently unconditionally allocate a ModuleState for all modules, even though nothing uses it yet)
  2. Very likely, changes to the public state API in PyModuleMethods needs work
    • Is &mut self valid? I see every other method -- including those that do mutation (add_*) -- take &self which makes me think this trait is expected handle synchronization internally
    • Notably absent is a method to delete a state type once initialized
    • I'm unsure what happens if two python thread states attempt to access a module's state concurrently... does something (GIL?) synchronize this for us?

bazaah avatar Nov 04 '25 18:11 bazaah

got an ICE compiling chrono?

https://github.com/PyO3/pyo3/actions/runs/19233273660/job/54976905090?pr=5600

   error: internal compiler error: compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs:777:17: `resolve_ty_and_res_fully_qualified_call` called on `LangItem`
  
  
  thread 'rustc' (25471) panicked at compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs:777:17:
  Box<dyn Any>
  stack backtrace:
     0: std::panicking::begin_panic::<rustc_errors::ExplicitBug>
     1: <rustc_errors::diagnostic::BugAbort as rustc_errors::diagnostic::EmissionGuarantee>::emit_producing_guarantee
     2: rustc_middle::util::bug::opt_span_bug_fmt::<rustc_span::span_encoding::Span>::{closure#0}
     3: rustc_middle::ty::context::tls::with_opt::<rustc_middle::util::bug::opt_span_bug_fmt<rustc_span::span_encoding::Span>::{closure#0}, !>::{closure#0}
     4: rustc_middle::ty::context::tls::with_context_opt::<rustc_middle::ty::context::tls::with_opt<rustc_middle::util::bug::opt_span_bug_fmt<rustc_span::span_encoding::Span>::{closure#0}, !>::{closure#0}, !>
     5: rustc_middle::util::bug::bug_fmt
     6: <rustc_hir_typeck::fn_ctxt::FnCtxt>::resolve_ty_and_res_fully_qualified_call
     7: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_path
     8: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_kind
     9: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_with_expectation_and_args
    10: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_kind
    11: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_with_expectation_and_args
    12: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_kind
    13: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_with_expectation_and_args
    14: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_block
    15: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_expr_with_expectation_and_args
    16: <rustc_hir_typeck::fn_ctxt::FnCtxt>::check_return_or_body_tail
    17: rustc_hir_typeck::check::check_fn
    18: rustc_hir_typeck::typeck_with_inspect::{closure#0}
    19: rustc_hir_typeck::typeck
        [... omitted 1 frame ...]
    20: <rustc_middle::ty::context::TyCtxt>::par_hir_body_owners::<rustc_hir_analysis::check_crate::{closure#2}>::{closure#0}
    21: rustc_hir_analysis::check_crate
    22: rustc_interface::passes::run_required_analyses
    23: rustc_interface::passes::analysis
        [... omitted 1 frame ...]
    24: rustc_interface::passes::create_and_enter_global_ctxt::<core::option::Option<rustc_interface::queries::Linker>, rustc_driver_impl::run_compiler::{closure#0}::{closure#2}>
    25: rustc_interface::interface::run_compiler::<(), rustc_driver_impl::run_compiler::{closure#0}>::{closure#1}
  note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.
  
  note: we would appreciate a bug report: https://github.com/rust-lang/rust/issues/new?labels=C-bug%2C+I-ICE%2C+T-compiler&template=ice.md
  
  note: rustc 1.91.0 (f8297e351 2025-10-28) running on aarch64-apple-darwin
  
  note: compiler flags: --crate-type lib -C embed-bitcode=no -C debuginfo=2 -C split-debuginfo=unpacked -C instrument-coverage
  
  note: some of the compiler flags provided by cargo are hidden
  
  query stack during panic:
  #0 [typeck] type-checking `format::<impl at /Users/runner/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/chrono-0.4.42/src/format/mod.rs:440:1: 440:33>::fmt`
  #1 [analysis] running analysis passes on this crate
  end of query stack
  error: could not compile `chrono` (lib)

That's new

EDIT: backlink https://github.com/rust-lang/rust/issues/148787

bazaah avatar Nov 10 '25 13:11 bazaah

CodSpeed Performance Report

Merging #5600 will not alter performance

Comparing bazaah:feat/module-state-core (93eb87e) with main (5fcc5d7)

Summary

✅ 98 untouched

codspeed-hq[bot] avatar Nov 14 '25 12:11 codspeed-hq[bot]

Thanks, this seems like a reasonable step forward. The StateCapsule type is obviously completely empty, I'd prefer we had some tangible substance before merging.

I don't yet see the full picture of how they state type gets passed around various parts of the PyO3 user code, it would be helpful to understand the current direction and proposed APIs for users to write / read module state.

Okay, can do. I'm currently using a typemap implementation that is similar to https://crates.io/crates/type-map but with more flexible bounds (the ability to have a Send + !Sync bound) that I wrote for some of my internal stuff. I'm not exactly sure if it's the right fit yet because I still feel unsure about what memory characteristics pymodule states need (and if we want to support the ability to store arbitrary user types), and I haven't really worked out where is the best place to expose the user facing APIs. As mentioned in https://github.com/PyO3/pyo3/issues/2274#issuecomment-3546879410, I currently do it off PyModuleMethods via &mut self -- which is conspicuously different to the &self receiver all of the other trait methods use, so I think I'm probably misunderstanding the characteristics of the trait.

Let me work on pushing up the typemap code plus what I have for the user API and we can talk about it from there

bazaah avatar Nov 18 '25 14:11 bazaah

I've rebased, added the real statemap (src/internal/typemap.rs) and addressed most of the previous feedback.

I think the next steps here are to have a discussion around the public API to module state, because I'm pretty sure it's wrong-ish. Maybe we can feature-gate it? I'm not sure how seriously pyo3 takes API stability, yet.

bazaah avatar Nov 20 '25 19:11 bazaah

For module traverse / clear - probably the discussion in #5663 is relevant (could be derived on the module state type).

davidhewitt avatar Dec 07 '25 10:12 davidhewitt