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

bazaah · 2025-11-04T18:50:24Z

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...

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

Left for future work:

Performance optimizations (we currently unconditionally allocate a ModuleState for all modules, even though nothing uses it yet)
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?

src/impl_/pymodule.rs

bazaah · 2025-11-10T13:39:23Z

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 rust-lang/rust#148787

codspeed-hq · 2025-11-14T12:13:04Z

CodSpeed Performance Report

Merging #5600 will not alter performance

_{Comparing bazaah:feat/module-state-core (4075435) with main (d8e9a38)}

Summary

✅ 98 untouched

davidhewitt

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.

davidhewitt · 2025-11-18T12:53:08Z

src/impl_/pymodule_state.rs

+    /// Calling this function multiple times on a single ModuleState is a noop,
+    /// beyond the first
+    fn drop_impl(&mut self) {
+        if let Some(ptr) = self.inner.take().map(|state| state.as_ptr()) {


Rather than Option<NonNull<_>> probably can use ManuallyDrop here.

The invariant "should not be accessed again" makes this sound like unsafe fn is appropriate here.

To be clear, this is a logic error rather than a safety one -- dropping the state without going through the m_free callback to drop the owning module can't result in memory safety issues, but will make our (and consumer) functions that expect a state to exist when it doesn't do the wrong thing. I'm happy to make it unsafe though it really shouldn't ever be called outside exactly the two cases in the code (Drop::drop & m_free callbacks)

src/impl_/pymodule_state.rs

src/impl_/pymodule.rs

bazaah · 2025-11-18T14:12:55Z

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 #2274 (comment), 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

This commit adds the core ModuleState & StateCapsule structs for interop with python's ModuleDef m_size field. We choose to keep the in-interpreter size small, only being a NonNull ptr to the actual (allocated by Rust) state. This is suboptimal from a pointer chasing perspective, because we'll eventually have interpreter->StateCapsule->TypeMap<T> dereference chains, but it makes it easier to implement and reason about.

This commit sets ModuleDef.m_size to size_of::<ModuleState>, and adds the relevant callback to ModuleDef.m_free. We also add C function defs for Py_mod_exec and ModuleDef.m_free callbacks. These will be used to initialize and free ModuleState in later patches. We do not yet set an initializer, as these are passed in m_slots, which are set from the `#[pymodule]` impl in pyo3-macros-backend.

Asserting we initialize ModuleState correctly, that it is accessible from both module_exec calls -- that is, during module initialization -- and from after the module is loaded.

This is ported from some of my internal code, with adjustments for edition = 2024. It was originally forked from rust-typemap but has diverged significantly, and been updated for newer rust versions. I freely relicense it under PyO3's licensing terms.

Within we switch over to the real typemap backed StateMap implementation. We also add a marker trait, ModuleStateType which guarantees the properties we expect from stored StateMap types. Namely: Clone + Send (not Sync), though I'm not sure these are the right guarantees for python's memory model. Send is a requirement, as python can move PyModules between threads at will, and Clone makes dealing with ModuleState much easier longer term (speaking from experience with typemaps). I'm not sure if Sync should be a requirement too, as it's not clear to me if PyModules can be concurrently accessed between attached python thread states. I'm worried the answer might change between free-threaded and "normal" python builds.

…tors

These commit adds three methods to the external interface for PyModules, exposing a limited ability to interact with the per-module state area. - state_ref - state_mut - state_or_init We take a hands-off approach to what the concrete "shape" of this state is, beyond requiring it is compatible with the marker ModuleStateType trait. Modules may store as many types as needed, and all the normal rust visibility conventions apply, allowing them to guard type invariants, or prevent external access entirely.

So we can return a real error if our state ptr is null. The upstream cpython builtins just assert non-null, so this is probably fine References: https://github.com/python/cpython/blob/v3.13.0/Modules/_sqlite/module.h#L80-L81 References: https://github.com/python/cpython/blob/v3.13.0/Modules/_io/_iomodule.h#L172-L173 References: https://github.com/python/cpython/blob/v3.13.0/Modules/_blake2/blake2module.c#L33-L34

bazaah · 2025-11-20T19:08:41Z

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 force-pushed the feat/module-state-core branch 3 times, most recently from 90fec38 to d4360ed Compare November 10, 2025 13:29

bazaah marked this pull request as ready for review November 10, 2025 13:29

bazaah commented Nov 10, 2025

View reviewed changes

src/impl_/pymodule.rs Outdated Show resolved Hide resolved

bazaah mentioned this pull request Nov 10, 2025

ICE while compiling chrono v0.4.42 (rust 1.91.0, aarch64-apple-darwin) in PyO3 CI rust-lang/rust#148787

Open

bazaah force-pushed the feat/module-state-core branch from 715f9c9 to 8c05240 Compare November 14, 2025 11:59

bazaah mentioned this pull request Nov 17, 2025

PEP 489 Compatibility #2274

Open

davidhewitt mentioned this pull request Nov 18, 2025

Implement #[init] method attribute in #[pymethods] #4951

Open

davidhewitt reviewed Nov 18, 2025

View reviewed changes

bazaah added 4 commits November 20, 2025 11:54

src/impl_: add test module_state_init

87554f7

Asserting we initialize ModuleState correctly, that it is accessible from both module_exec calls -- that is, during module initialization -- and from after the module is loaded.

pyo3-macros-backend: add Py_mod_exec slot for initializing ModuleState

c92eeb3

bazaah force-pushed the feat/module-state-core branch from 8c05240 to c09543c Compare November 20, 2025 11:54

bazaah force-pushed the feat/module-state-core branch from 2613c9a to 6f5b849 Compare November 20, 2025 12:45

bazaah added 8 commits November 20, 2025 18:19

src/impl_: add getters for ModuleState->StateMap, from_bound construc…

0a6d91d

…tors

src/impl_: update test to use PyModuleMethods, assert on dummy UserState

4b36061

src/impl_: inline size_of<ModuleState>

94f472f

src/impl_: make ModuleState::drop_impl unsafe, document usage

a17624e

newsfragments: add PR description

4075435

bazaah force-pushed the feat/module-state-core branch from 6f5b849 to 4075435 Compare November 20, 2025 18:20

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

bazaah commented Nov 4, 2025 •

edited

Loading

Uh oh!

Uh oh!

bazaah commented Nov 10, 2025 •

edited

Loading

Uh oh!

codspeed-hq bot commented Nov 14, 2025 •

edited

Loading

Uh oh!

davidhewitt left a comment

Uh oh!

davidhewitt Nov 18, 2025

Uh oh!

bazaah Nov 18, 2025

Uh oh!

Uh oh!

Uh oh!

bazaah commented Nov 18, 2025

Uh oh!

bazaah commented Nov 20, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

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

Are you sure you want to change the base?

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

Conversation

bazaah commented Nov 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

bazaah commented Nov 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codspeed-hq bot commented Nov 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

CodSpeed Performance Report

Merging #5600 will not alter performance

Summary

Uh oh!

davidhewitt left a comment

Choose a reason for hiding this comment

Uh oh!

davidhewitt Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

bazaah Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

bazaah commented Nov 18, 2025

Uh oh!

bazaah commented Nov 20, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

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

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

bazaah commented Nov 4, 2025 •

edited

Loading

bazaah commented Nov 10, 2025 •

edited

Loading

codspeed-hq bot commented Nov 14, 2025 •

edited

Loading