Add module state support - #6245
Draft
MatthieuDartiailh wants to merge 24 commits into
Draft
Conversation
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. (cherry picked from commit 33e8324)
This commit sets ModuleDef.m_size to size_of::<ModuleState>, and adds the relevant callback to ModuleDef.m_free.
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. (cherry picked from commit fabaf73)
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. (cherry picked from commit 4f0c205)
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. (cherry picked from commit 1de1b3f)
…tors (cherry picked from commit 0b2ad5d)
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. (cherry picked from commit 2e7b1f2)
(cherry picked from commit eaf9827)
(cherry picked from commit 6806ae8)
(cherry picked from commit 7be3e8a)
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 (cherry picked from commit 78dd17f)
(cherry picked from commit 5987894)
… with pymodule The module state is allocated only if a module state is requested, module init function must return a state instance.
… footgunny impl from PyAny Also add the machinery to allow attaching a module to a type This is required to be able to access the module state from the type at low cost
MatthieuDartiailh
force-pushed
the
module-state-support
branch
from
July 27, 2026 19:04
c5cb30a to
38c18c2
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This builds upon #5600 but takes a slightly different approach to managing and accessing the module state:
pymoduleusing a newstatekwargs or declared inside a mod using#[pymodule_state]pymoduleis applied#[pymodule_init]function inside a mod decorated with#[pymodule]#[pymodule_traverse]and#[pymodule_clear]As a consequence there is single module state and there is no need for a type map to store it a
Box<dyn Any>is sufficient.The state can be accessed either through a module object using
module_stateormodule_state_mutwhich is unsafe (since we cannot guarantee the unicity), locking the module using a critical sction can make it safe, however in practice I would expect Mutex to be used on the state to provide a finer granularity.It can also be accessed from a type object using
module_state(no mutable variant since the user does not get the module to lock) or from the Python token throughtype_module_statewhich takes 2 type parameters (one describes the type, one the state).For this to work the type needs to be created using PyType_FromModuleAndSpec. For this I added
add_class_with_moduleto the Module object, which will properly create the class and error if the class was already created. (If required I could probably split this into a separate PR)An example is included (compiles and passes tests).
This is draft because: