diff --git a/Cargo.toml b/Cargo.toml index ae72c983804..d61757e90a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -99,6 +99,9 @@ pyo3-build-config = { path = "pyo3-build-config", version = "=0.29.0" } [features] default = ["macros"] +# Enables support defining Python module state backed by a Rust structure. +experimental-module-state = ["macros", "pyo3-macros/experimental-module-state"] + # Enables support for `async fn` for `#[pyfunction]` and `#[pymethods]`. experimental-async = ["macros", "pyo3-macros/experimental-async"] diff --git a/examples/module_state/.template/.keep b/examples/module_state/.template/.keep new file mode 100644 index 00000000000..6ca24f1959b --- /dev/null +++ b/examples/module_state/.template/.keep @@ -0,0 +1 @@ +# This directory is used by cargo-generate as a template diff --git a/examples/module_state/Cargo.toml b/examples/module_state/Cargo.toml new file mode 100644 index 00000000000..fbbf58f3e65 --- /dev/null +++ b/examples/module_state/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "module_state_example" +version = "0.1.0" +edition = "2021" +rust-version = "1.83" + +[lib] +name = "module_state" +crate-type = ["cdylib"] + +[dependencies] +pyo3 = { path = "../../", features = ["experimental-module-state"] } + +[workspace] diff --git a/examples/module_state/MANIFEST.in b/examples/module_state/MANIFEST.in new file mode 100644 index 00000000000..d4437e7fe07 --- /dev/null +++ b/examples/module_state/MANIFEST.in @@ -0,0 +1,3 @@ +recursive-include src * +recursive-include tests * +include README.md diff --git a/examples/module_state/README.md b/examples/module_state/README.md new file mode 100644 index 00000000000..c3c10a436a2 --- /dev/null +++ b/examples/module_state/README.md @@ -0,0 +1,308 @@ +# Module State Example + +This example demonstrates the **experimental module state API** that will be available in PyO3 once the `experimental-module-state` feature is stabilized. + +## Overview + +The module state API allows you to: + +- Define per-module state that persists across function calls +- Initialize state once during module import +- Access state from functions and methods safely +- Share data across the module (like configuration, caches, counters, etc.) + +## Key API Features + +### 1. **State Struct Definition** + +```rust +#[pymodule_state] +struct ModuleState { + counter: i32, + config: String, + data: Arc, +} +``` + +The `#[pymodule_state]` marker tells PyO3 that this struct represents the module's state. + +### 2. **Module Declaration with State** + +```rust +#[pymodule(state = ModuleState)] +fn module_state(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult { + m.add_function(wrap_pyfunction!(get_counter, m)?)?; + // ... add more functions/classes + ModuleState::new() // Initialize and return state +} +``` + +The `state = ModuleState` parameter tells PyO3 to manage `ModuleState` for this module. + +### 3. **Accessing State from Functions with `pass_module`** + +Functions that need to access module state must use the `pass_module` attribute: + +```rust +/// Functions need pass_module to receive the module parameter +#[pyfunction(pass_module)] +fn get_counter(m: &Bound<'_, PyModule>) -> PyResult { + // Safe, immutable access returns Option<&T> + if let Some(state) = m.module_state::() { + Ok(state.counter) + } else { + Err(PyErr::new::( + "Module state not available" + )) + } +} +``` + +**Key Points**: + +- `#[pyfunction(pass_module)]` attribute tells PyO3 to inject the module as a parameter +- The module parameter comes before regular arguments +- Returns `Option<&T>` - handles missing state gracefully + +### 4. **Mutating State from Functions** + +```rust +#[pyfunction(pass_module)] +fn increment(m: &Bound<'_, PyModule>) -> PyResult { + // Unsafe, mutable access returns Option<&mut T> + unsafe { + if let Some(state) = m.module_state_mut::() { + state.counter += 1; + Ok(state.counter) + } else { + Err(...) + } + } +} +``` + +**Safety**: Requires explicit `unsafe` block. +Safe because: + +- You're under the GIL (PyModule access implies it) +- You have sole mutable access to the module's state +- Generally only used during initialization + +### 5. **Accessing State from Classes with `PyAnyMethods` (Optimized)** + +Classes can access their defining module's state via `PyAnyMethods` - available on any object without reference counting overhead: + +```rust +#[pymethods] +impl Counter { + fn get_module_state_value(slf: &Bound<'_, Self>) -> PyResult { + // Bound<'_, Counter> coerces to Bound<'_, PyAny> + // type_module_state() uses Py_TYPE() internally - no incref/decref! + if let Some(state) = slf.type_module_state::() { + Ok(state.counter) + } else { + Err(PyErr::new::( + "Module state not available" + )) + } + } + + fn update_config(slf: &Bound<'_, Self>, new_config: String) -> PyResult<()> { + // Mutable access from class method + unsafe { + if let Some(state) = slf.type_module_state_mut::() { + state.config = new_config; + Ok(()) + } else { + Err(PyErr::new::( + "Module state not available" + )) + } + } + } +} +``` + +**Key Points**: + +- `type_module_state::()` - safe immutable access via Py_TYPE() + PyType_GetModuleState +- `type_module_state_mut::()` - mutable access (requires unsafe) +- **No reference counting overhead** - uses `Py_TYPE()` directly instead of creating a `Bound<'_, PyType>` +- Works on any `Bound<'_, PyAny>`, including instances, not just types +- No need to call `get_type()` first + +### 6. **Subinterpreter Isolation** + +Each subinterpreter gets its own isolated state instance. +This is handled automatically by PyO3's module lifecycle system. + +## What This Example Shows + +### Immutable State Access + +- `get_counter()` - read-only counter with `pass_module` +- `get_config()` - read-only configuration with `pass_module` +- `get_data()` - access shared Arc data with `pass_module` + +### Mutable State Access + +- `increment()` - modify counter with `pass_module` +- `set_config()` - update configuration with `pass_module` + +### Class Integration with PyAnyMethods + +- `Counter` class uses `type_module_state::()` to access module state +- `get_module_state_value()` - immutable access from class method +- `update_config()` - mutable access from class method +- No need to import module or pass it around +- **Optimized**: Uses Py_TYPE() directly - no reference counting overhead + +## Running the Example (Once Implemented) + +```bash +# Build the module +cargo build --example module_state --features experimental-module-state + +# In Python: +import module_state + +# Read state via functions with pass_module +print(module_state.get_counter()) # 0 + +# Modify state +print(module_state.increment()) # 1 +print(module_state.increment()) # 2 + +# Configure state +module_state.set_config("custom") +print(module_state.get_config()) # "custom" + +# Access from class via PyTypeMethods +counter = module_state.Counter("test") +print(counter.get_module_state_value()) # 2 +counter.update_config("from_class") +print(module_state.get_config()) # "from_class" +``` + +## Design Rationale + +### Single State Type Per Module + +Only **one** state type per module (not multiple types). +This: + +- Simplifies the API and implementation +- Makes the state type known at compile-time +- Enables macro validation of init function return types +- Matches real-world usage patterns + +### `pass_module` for Functions + +The `pass_module` attribute on `#[pyfunction]`: + +- Explicitly declares that a function needs the module +- Makes it clear in the code where module state is accessed +- PyO3 injects the module automatically before other arguments +- Better than implicit module parameter passing + +### PyAnyMethods for Classes + +Classes can access module state from any instance efficiently: + +- `type_module_state::()` uses Py_TYPE() + PyType_GetModuleState internally +- **No reference counting overhead** - works directly with borrowed type pointer +- Available on any `Bound<'_, PyAny>` (includes class instances) +- Works automatically for any class defined in a module with state +- No need to import or pass the module around in class methods +- More general than type-only API (works on instances too) + +### Option-Based Access + +`module_state()` and `type_module_state()` return `Option<&T>`: + +- Safe handling of initialization failures +- Clear error path in code +- No surprising panics from type mismatches + +### Explicit Unsafe for Mutation + +Mutable access requires `unsafe`: + +- Signals that mutable access needs care +- Documents the GIL assumption +- Encourages thinking about thread safety + +## Comparison with PR #5600 + +This example shows the **redesigned API**, which differs from PR #5600: + +| Aspect | PR #5600 | New Design | +|--------|----------|-----------| +| State Storage | `TypeMap` (multiple types) | `Box` (single type) | +| PyModule API | `state_ref()`, `state_mut()` | `module_state()`, `module_state_mut()` | +| PyAny API | None | `type_module_state()`, `type_module_state_mut()` (via PyAnyMethods) | +| Function Access | No mechanism | `#[pyfunction(pass_module)]` | +| Class Access | No mechanism | Direct via `PyAnyMethods` on instance | +| Error Handling | `PyResult` | `Option` (None for missing) | +| Feature Gate | No | `experimental-module-state` | +| Type Validation | Runtime downcast | Compile-time macro validation | + +## Key Differences from PR #5600 + +### State Storage + +- **PR #5600**: Uses `TypeMap` to store multiple state types (complex) +- **New Design**: Single `Box` per module (simple, matches real use cases) + +### Function Access Pattern + +- **PR #5600**: No standard pattern for function access +- **New Design**: Use `#[pyfunction(pass_module)]` to receive module and call `m.module_state::()` + +### Class Access Pattern + +- **PR #5600**: No standard pattern for class access +- **New Design**: Use `PyAnyMethods` on any object instance to call `obj.type_module_state::()` via Py_TYPE() + PyType_GetModuleState (no refcount overhead) + +### API Method Names + +- **PR #5600**: `state_ref()`, `state_mut()` (generic names) +- **New Design**: + - `module_state()`, `module_state_mut()` on PyModule (domain-specific) + - `type_module_state()`, `type_module_state_mut()` on PyAny (optimized, no refcount) + +## Implementation Status + +This example serves as the **guiding specification** for implementation. +Phases: + +1. ✅ Phase 2.0: Fix proc macro slot counting bug +2. 🟡 Phase 2.1-2.7: Implement the API (in progress) + - Phase 2.1: `#[pymodule_state]` macro + - Phase 2.2: Parser extension + auto-detection + - Phase 2.3: Function-level module handling + - Phase 2.4: Return type validation + - Phase 2.5: State initialization code generation + - Phase 2.6: PyModule API methods + - **Phase 2.6.3: PyAnyMethods for optimized class access** (NEW - uses Py_TYPE, no refcount) + - Phase 2.7: Cleanup +3. ❌ Phase 3: Full testing and validation + +See [PHASE2_DETAILED_IMPLEMENTATION_PLAN.md](../../PHASE2_DETAILED_IMPLEMENTATION_PLAN.md) for details. + +## Next Steps + +To make this example compile: + +1. Implement `#[pymodule_state]` marker macro +2. Extend `#[pymodule(state = ...)]` parser +3. Implement state type auto-detection +4. Add `pass_module` support to `#[pyfunction]` +5. Add `module_state()` and `module_state_mut()` methods +6. Add PyType_GetModuleState FFI binding +7. Implement `PyAnyMethods` with `type_module_state()` and `type_module_state_mut()` + - Uses Py_TYPE() directly for zero-cost access + - Available on any `Bound<'_, PyAny>` +8. Generate state initialization code + +Each step is tracked in the implementation plan. diff --git a/examples/module_state/cargo-generate.toml b/examples/module_state/cargo-generate.toml new file mode 100644 index 00000000000..c9249729b98 --- /dev/null +++ b/examples/module_state/cargo-generate.toml @@ -0,0 +1,15 @@ +[template] +# This file is used by cargo-generate if someone wants to use this +# as a template for creating new module state projects +exclude = [".git", "target", ".template"] + +[placeholders.project_name] +type = "string" +prompt = "What is the name of your project?" +regex = "^[a-zA-Z][a-zA-Z0-9_-]*$" + +[placeholders.module_name] +type = "string" +prompt = "What would you like your module to be called in Python?" +regex = "^[a-z][a-z0-9_]*$" +default = "my_module" diff --git a/examples/module_state/noxfile.py b/examples/module_state/noxfile.py new file mode 100644 index 00000000000..ba54d0582a2 --- /dev/null +++ b/examples/module_state/noxfile.py @@ -0,0 +1,23 @@ +import nox + + +@nox.session +def test(session): + """Run tests for module state example (once implemented).""" + session.install("pytest") + session.run("pytest", "tests/", "-v") + + +@nox.session +def build(session): + """Build the module.""" + session.install("maturin") + session.run("maturin", "develop") + + +@nox.session +def dev(session): + """Development session with build and test.""" + session.run("maturin", "develop", external=True) + session.install("pytest") + session.run("pytest", "tests/", "-v") diff --git a/examples/module_state/pyproject.toml b/examples/module_state/pyproject.toml new file mode 100644 index 00000000000..edfda2fe39e --- /dev/null +++ b/examples/module_state/pyproject.toml @@ -0,0 +1,7 @@ +[build-system] +requires = ["maturin>=0.14,<0.15"] +build-backend = "maturin" + +[project] +name = "module-state-example" +version = "0.1.0" diff --git a/examples/module_state/src/lib.rs b/examples/module_state/src/lib.rs new file mode 100644 index 00000000000..c57ede70282 --- /dev/null +++ b/examples/module_state/src/lib.rs @@ -0,0 +1,156 @@ +//! Example: Module State API (Experimental/Future) +//! +//! This example demonstrates the declarative module state API that will be available +//! once the `experimental-module-state` feature is stabilized. +//! +//! The key features shown here: +//! - `#[pymodule_state]` marker on a state struct +//! - `#[pymodule(state = ModuleState)]` declarative state binding +//! - `module_state()` and `module_state_mut()` API for accessing state +//! - Type-safe state storage per module + +use pyo3::prelude::*; +use std::sync::Mutex; + +/// Module state struct marked with `#[pymodule_state]` +/// +/// This struct will be automatically instantiated during module initialization +/// and made available to all functions in the module via `module_state()`. +struct ModuleState { + /// Counter that increments with each call to `increment()` + counter: Mutex, + /// Configuration string set at init time + config: String, +} + +impl ModuleState { + /// Initialize the module state + /// + /// This is called once when the module is first imported. + /// Any initialization errors here are propagated to Python. + fn new() -> PyResult { + println!("Initializing module state..."); + Ok(ModuleState { + counter: Mutex::new(0), + config: "initialized".to_string(), + }) + } +} + +/// Get the current counter value from module state +/// +/// Functions that need to access module state must use `pass_module` to receive +/// the module as a parameter. +#[pyfunction(pass_module)] +fn get_counter(m: &Bound<'_, PyModule>) -> PyResult { + // Safe API: returns Option<&T>, handles missing/wrong type + if let Some(state) = m.module_state::() { + Ok(*state.counter.lock().unwrap()) + } else { + Err(PyErr::new::( + "Module state not available", + )) + } +} + +/// Increment the counter in module state +/// +/// Note: This requires `unsafe` access to get a mutable reference to state. +/// Generally safe during initialization or under the GIL. +#[pyfunction(pass_module)] +fn increment_counter(m: &Bound<'_, PyModule>) -> PyResult { + let new_value = { + if let Some(state) = m.module_state::() { + *state.counter.lock().unwrap() += 1; + *state.counter.lock().unwrap() + } else { + return Err(PyErr::new::( + "Module state not available", + )); + } + }; + Ok(new_value) +} + +/// Get the configuration from module state +#[pyfunction(pass_module)] +fn get_config(m: &Bound<'_, PyModule>) -> PyResult { + if let Some(state) = m.module_state::() { + Ok(state.config.clone()) + } else { + Err(PyErr::new::( + "Module state not available", + )) + } +} + +/// Example of accessing state from a class method +/// +/// Classes defined in a module with state can access that state via Python token. +#[pyclass(module = "module_state")] +struct Counter { + name: String, +} + +#[pymethods] +impl Counter { + #[new] + fn new(name: String) -> Self { + Counter { name } + } + + /// Access module state from a class method using Python token + fn get_shared_counter_value(slf: &Bound<'_, Self>) -> PyResult { + if let Some(state) = slf.py().type_module_state::()? { + Ok(*state.counter.lock().unwrap()) + } else { + Err(PyErr::new::( + "Module state not available", + )) + } + } + + /// Update module state from a class method + fn increment_shared_counter(slf: &Bound<'_, Self>) -> PyResult { + if let Some(state) = slf.py().type_module_state::()? { + *state.counter.lock().unwrap() += 1; + Ok(*state.counter.lock().unwrap()) + } else { + Err(PyErr::new::( + "Module state not available", + )) + } + } +} + +/// Module definition using declarative state API +#[pymodule(state = ModuleState)] +fn module_state(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult { + // Add functions to the module + m.add_function(wrap_pyfunction!(get_counter, m)?)?; + m.add_function(wrap_pyfunction!(increment_counter, m)?)?; + m.add_function(wrap_pyfunction!(get_config, m)?)?; + + // Add classes to the module + m.add_class_with_module::()?; + + // Add module docstring + m.add( + "__doc__", + "Module State Example\n\ + \n\ + This module demonstrates the experimental module state API.\n\ + It maintains per-module state that persists across function calls.\n\ + \n\ + Functions:\n\ + - get_counter(): Get the current counter value from module state\n\ + - increment_counter(): Increment the counter and return new value\n\ + - get_config(): Get the current configuration string\n\ + \n\ + Classes:\n\ + - Counter: Example class with method accessing module state\ + ", + )?; + + ModuleState::new() +} diff --git a/examples/module_state/tests/test_module.py b/examples/module_state/tests/test_module.py new file mode 100644 index 00000000000..00e86d2c39d --- /dev/null +++ b/examples/module_state/tests/test_module.py @@ -0,0 +1,56 @@ +""" +Test file for module state example. + +This demonstrates how the module state API would be used in Python tests. +It won't run until the feature is implemented. +""" + +import pytest + + +class TestModuleState: + """Tests for module state functionality.""" + + def test_counter_initial_value(self): + """Module state counter starts at 0.""" + import module_state + + assert module_state.get_counter() == 0 + + def test_counter_increment(self): + """Counter increments correctly.""" + import module_state + + assert module_state.increment_counter() == 1 + assert module_state.increment_counter() == 2 + assert module_state.get_counter() == 2 + + def test_counter_persists_across_calls(self): + """Counter value persists across multiple function calls.""" + import module_state + + # Start fresh for determinism (in real tests, use fixtures) + initial = module_state.get_counter() + module_state.increment_counter() + module_state.increment_counter() + assert module_state.get_counter() == initial + 2 + + def test_config_get_set(self): + """Configuration can be read and written.""" + import module_state + + assert module_state.get_config() == "initialized" + + def test_counter_class_access(self): + """Counter class can access module state.""" + import module_state + + counter = module_state.Counter("test_counter") + value = counter.get_shared_counter_value() + assert isinstance(value, int) + # Value should match the module-level counter + assert value == module_state.get_counter() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/newsfragments/6245.added.md b/newsfragments/6245.added.md new file mode 100644 index 00000000000..9402fd3af26 --- /dev/null +++ b/newsfragments/6245.added.md @@ -0,0 +1 @@ +Add support for attaching a module state when creating a pymodule diff --git a/pyo3-macros-backend/Cargo.toml b/pyo3-macros-backend/Cargo.toml index add4398df10..9fc0e3fd761 100644 --- a/pyo3-macros-backend/Cargo.toml +++ b/pyo3-macros-backend/Cargo.toml @@ -30,3 +30,4 @@ workspace = true [features] experimental-async = [] experimental-inspect = [] +experimental-module-state = [] diff --git a/pyo3-macros-backend/src/attributes.rs b/pyo3-macros-backend/src/attributes.rs index 9894c463628..87dd42d3801 100644 --- a/pyo3-macros-backend/src/attributes.rs +++ b/pyo3-macros-backend/src/attributes.rs @@ -56,6 +56,8 @@ pub mod kw { syn::custom_keyword!(category); syn::custom_keyword!(from_py_object); syn::custom_keyword!(skip_from_py_object); + #[cfg(feature = "experimental-module-state")] + syn::custom_keyword!(state); } fn take_int(read: &mut &str, tracker: &mut usize) -> String { @@ -349,6 +351,8 @@ pub type TextSignatureAttribute = KeywordAttribute; pub type SubmoduleAttribute = kw::submodule; pub type GILUsedAttribute = KeywordAttribute; +#[cfg(feature = "experimental-module-state")] +pub type StateAttribute = KeywordAttribute; impl Parse for KeywordAttribute { fn parse(input: ParseStream<'_>) -> Result { diff --git a/pyo3-macros-backend/src/module.rs b/pyo3-macros-backend/src/module.rs index 44164ef48b9..7a0a2d693ae 100644 --- a/pyo3-macros-backend/src/module.rs +++ b/pyo3-macros-backend/src/module.rs @@ -1,5 +1,7 @@ //! Code generation for the function that initializes a python module and adds classes and function. +#[cfg(feature = "experimental-module-state")] +use crate::attributes::StateAttribute; #[cfg(feature = "experimental-inspect")] use crate::introspection::{ attribute_introspection_code, introspection_id_const, module_introspection_code, @@ -38,6 +40,8 @@ pub struct PyModuleOptions { module: Option, submodule: Option, gil_used: Option, + #[cfg(feature = "experimental-module-state")] + state: Option, } impl Parse for PyModuleOptions { @@ -86,6 +90,10 @@ impl PyModuleOptions { PyModulePyO3Option::GILUsed(gil_used) => { set_option!(gil_used) } + #[cfg(feature = "experimental-module-state")] + PyModulePyO3Option::State(state) => { + set_option!(state) + } } Ok(()) @@ -163,7 +171,22 @@ pub fn pymodule_module_impl( Ok(()) } - let mut pymodule_init = None; + #[cfg(feature = "experimental-module-state")] + let mut state_type: Option = None; + #[cfg(feature = "experimental-module-state")] + if let Some(explicit_state) = &options.state { + state_type = Some(explicit_state.value.clone()); + } + + let mut pymodule_init: Option = None; + #[cfg(feature = "experimental-module-state")] + let mut pymodule_traverse: Option = None; + #[cfg(not(feature = "experimental-module-state"))] + let pymodule_traverse: Option = None; + #[cfg(feature = "experimental-module-state")] + let mut pymodule_clear: Option = None; + #[cfg(not(feature = "experimental-module-state"))] + let pymodule_clear: Option = None; let mut module_consts = Vec::new(); let mut module_consts_cfg_attrs = Vec::new(); @@ -189,14 +212,37 @@ pub fn pymodule_module_impl( ); let is_pymodule_init = find_and_remove_attribute(&mut item_fn.attrs, "pymodule_init"); + #[cfg(feature = "experimental-module-state")] + let is_pymodule_traverse = + find_and_remove_attribute(&mut item_fn.attrs, "pymodule_traverse"); + #[cfg(feature = "experimental-module-state")] + let is_pymodule_clear = + find_and_remove_attribute(&mut item_fn.attrs, "pymodule_clear"); let ident = &item_fn.sig.ident; + #[cfg(feature = "experimental-module-state")] + if is_pymodule_traverse { + ensure_spanned!( + !has_attribute(&item_fn.attrs, "pyfunction"), + item_fn.span() => "`#[pyfunction]` cannot be used alongside `#[pymodule_traverse]`" + ); + ensure_spanned!(pymodule_traverse.is_none(), item_fn.span() => "only one `#[pymodule_traverse]` may be specified"); + pymodule_traverse = Some(quote! { #ident }); + } else if is_pymodule_clear { + ensure_spanned!( + !has_attribute(&item_fn.attrs, "pyfunction"), + item_fn.span() => "`#[pyfunction]` cannot be used alongside `#[pymodule_clear]`" + ); + ensure_spanned!(pymodule_clear.is_none(), item_fn.span() => "only one `#[pymodule_clear]` may be specified"); + pymodule_clear = Some(quote! { #ident }) + } + if is_pymodule_init { ensure_spanned!( !has_attribute(&item_fn.attrs, "pyfunction"), item_fn.span() => "`#[pyfunction]` cannot be used alongside `#[pymodule_init]`" ); ensure_spanned!(pymodule_init.is_none(), item_fn.span() => "only one `#[pymodule_init]` may be specified"); - pymodule_init = Some(quote! { #ident(module)?; }); + pymodule_init = Some(quote! { #ident(module) }); } else if has_attribute(&item_fn.attrs, "pyfunction") || has_attribute_with_namespace( &item_fn.attrs, @@ -240,6 +286,21 @@ pub fn pymodule_module_impl( set_module_attribute(&mut item_struct.attrs, &full_name); } } + #[cfg(feature = "experimental-module-state")] + if find_and_remove_attribute(&mut item_struct.attrs, "pymodule_state") { + if state_type.is_some() { + bail_spanned!(item_struct.span() => + "Multiple `#[pymodule_state]` structs found. Specify state type explicitly with `#[pymodule(state = ...)]`"); + } + state_type = Some( + syn::Path { + leading_colon: None, + segments: std::iter::once(syn::PathSegment { + ident: item_struct.ident.clone(), + arguments: syn::PathArguments::None, + }).collect(),} + ); + } } Item::Enum(item_enum) => { ensure_spanned!( @@ -393,6 +454,38 @@ pub fn pymodule_module_impl( let gil_used = options.gil_used.is_some_and(|op| op.value.value); + #[cfg(feature = "experimental-module-state")] + let allocate_state = { + // For mod declarations with state, require an init function + if let Some(state) = options.state.as_ref().map(|s| &s.value) { + if pymodule_init.is_none() { + return Err(syn::Error::new_spanned( + &ident, + "module with `state` attribute must have a `#[pymodule_init]` function that returns the state value" + )); + } + + pymodule_init = Some(quote! { + let state: #state = #pymodule_init?; + #pyo3_path::impl_::pymodule::pyo3_module_state_init(module, state) + }); + true + } else { + if pymodule_init.is_none() { + pymodule_init = Some(quote! { #pyo3_path::PyResult::Ok(()) }); + }; + + false + } + }; + #[cfg(not(feature = "experimental-module-state"))] + let allocate_state = { + if pymodule_init.is_none() { + pymodule_init = Some(quote! { #pyo3_path::PyResult::Ok(()) }); + }; + false + }; + let initialization = module_initialization( &full_name, &name, @@ -401,6 +494,9 @@ pub fn pymodule_module_impl( options.submodule.is_some(), gil_used, doc.as_ref(), + allocate_state, + pymodule_traverse, + pymodule_clear, )?; let module_consts_names = module_consts.iter().map(|i| i.unraw().to_string()); @@ -428,7 +524,6 @@ pub fn pymodule_module_impl( )* #pymodule_init - ::std::result::Result::Ok(()) } } )) @@ -453,6 +548,11 @@ pub fn pymodule_function_impl( let gil_used = options.gil_used.is_some_and(|op| op.value.value); + #[cfg(feature = "experimental-module-state")] + let allocate_state = options.state.is_some(); + #[cfg(not(feature = "experimental-module-state"))] + let allocate_state = false; + let initialization = module_initialization( &name.to_string(), &name, @@ -461,6 +561,9 @@ pub fn pymodule_function_impl( false, gil_used, doc.as_ref(), + allocate_state, + None, + None, )?; #[cfg(feature = "experimental-inspect")] @@ -486,6 +589,25 @@ pub fn pymodule_function_impl( } module_args.push(quote!(::std::convert::Into::into(module))); + #[cfg(feature = "experimental-module-state")] + let init_func = { + if let Some(state) = &options.state { + let state = &state.value; + quote! { + let state: #state = #ident(#(#module_args),*)?; + #pyo3_path::impl_::pymodule::pyo3_module_state_init(module, state) + } + } else { + quote! { + #ident(#(#module_args),*) + } + } + }; + #[cfg(not(feature = "experimental-module-state"))] + let init_func = quote! { + #ident(#(#module_args),*) + }; + Ok(quote! { #[doc(hidden)] #vis mod #ident { @@ -501,12 +623,13 @@ pub fn pymodule_function_impl( #[allow(unknown_lints, non_local_definitions)] impl #ident::ModuleExec { fn __pyo3_module_exec(module: &#pyo3_path::Bound<'_, #pyo3_path::types::PyModule>) -> #pyo3_path::PyResult<()> { - #ident(#(#module_args),*) + #init_func } } }) } +#[allow(clippy::too_many_arguments)] fn module_initialization( full_name: &str, name: &syn::Ident, @@ -515,6 +638,9 @@ fn module_initialization( is_submodule: bool, gil_used: bool, doc: Option<&PythonDoc>, + allocate_state: bool, + pymodule_traverse: Option, + pymodule_clear: Option, ) -> Result { let Ctx { pyo3_path, .. } = ctx; let pyinit_symbol = format!("PyInit_{name}"); @@ -526,6 +652,118 @@ fn module_initialization( c"".into_token_stream() }; + // Generate traverse and clear C callbacks if present + #[cfg(feature = "experimental-module-state")] + let (traverse_callback, clear_callback) = + if pymodule_traverse.is_some() || pymodule_clear.is_some() { + let traverse_code = if let Some(traverse_fn) = &pymodule_traverse { + quote! { + unsafe extern "C" fn __pyo3_module_traverse( + module: *mut #pyo3_path::ffi::PyObject, + visit: #pyo3_path::ffi::visitproc, + arg: *mut ::std::ffi::c_void, + ) -> ::std::ffi::c_int { + #pyo3_path::Python::with_gil(|py| { + let module_bound = #pyo3_path::Bound::new_borrowed(py, module); + let py_module = module_bound.cast_exact::<#pyo3_path::types::PyModule>() + .expect("module object is not PyModule"); + + match py_module.module_state::<::std::any::Any>() { + Ok(state_any) => { + let visit_fn = #pyo3_path::PyVisit { + visit, + arg, + _guard: ::std::marker::PhantomData, + }; + // Call user's traverse function with state and visit + match #traverse_fn(state_any, visit_fn) { + Ok(()) => 0, + Err(_) => -1, + } + } + Err(_) => -1, + } + }) + } + } + } else { + quote! {} + }; + + let clear_code = if let Some(clear_fn) = &pymodule_clear { + quote! { + unsafe extern "C" fn __pyo3_module_clear( + module: *mut #pyo3_path::ffi::PyObject, + ) -> ::std::ffi::c_int { + #pyo3_path::Python::with_gil(|py| { + let module_bound = #pyo3_path::Bound::new_borrowed(py, module); + let py_module = module_bound.cast_exact::<#pyo3_path::types::PyModule>() + .expect("module object is not PyModule"); + + match py_module.module_state_mut::<::std::any::Any>() { + Ok(state_any_mut) => { + // Call user's clear function with mutable state + match #clear_fn(state_any_mut) { + Ok(()) => 0, + Err(_) => -1, + } + } + Err(_) => -1, + } + }) + } + } + } else { + quote! {} + }; + + (traverse_code, clear_code) + } else { + (quote! {}, quote! {}) + }; + + #[cfg(not(feature = "experimental-module-state"))] + let (traverse_callback, clear_callback) = { + let (_, _) = (pymodule_traverse, pymodule_clear); + (quote! {}, quote! {}) + }; + + #[cfg(feature = "experimental-module-state")] + let (traverse_arg, clear_arg) = if pymodule_traverse.is_some() || pymodule_clear.is_some() { + let traverse = if pymodule_traverse.is_some() { + quote! { ::core::option::Option::Some(__pyo3_module_traverse) } + } else { + quote! { ::core::option::Option::None } + }; + let clear = if pymodule_clear.is_some() { + quote! { ::core::option::Option::Some(__pyo3_module_clear) } + } else { + quote! { ::core::option::Option::None } + }; + (traverse, clear) + } else { + ( + quote! { ::core::option::Option::None }, + quote! { ::core::option::Option::None }, + ) + }; + #[cfg(not(feature = "experimental-module-state"))] + let (traverse_arg, clear_arg) = ( + quote! { ::core::option::Option::None }, + quote! { ::core::option::Option::None }, + ); + + let mod_new = quote! { + #pyo3_path::impl_::pymodule::ModuleDef::new( + __PYO3_NAME, + #doc, + &SLOTS, + #allocate_state, + #traverse_arg, + #clear_arg, + ) + }; + let mut result = quote! { #[doc(hidden)] pub const __PYO3_NAME: &'static ::std::ffi::CStr = #pyo3_name; @@ -544,6 +782,9 @@ fn module_initialization( #pyo3_path::impl_::trampoline::module_exec(module, #module_exec) } + #traverse_callback + #clear_callback + // The full slots, used for the PyModExport initialization static SLOTS: impl_::PyModuleSlots = impl_::PyModuleSlotsBuilder::new() .with_mod_exec(__pyo3_module_exec) @@ -556,7 +797,7 @@ fn module_initialization( // Since the macros need to be written agnostic to the Python version // we need to explicitly pass the name and docstring for PyModuleDef // initialization. - impl_::ModuleDef::new(__PYO3_NAME, #doc, &SLOTS) + #mod_new }; }; if !is_submodule { @@ -720,11 +961,18 @@ enum PyModulePyO3Option { Name(NameAttribute), Module(ModuleAttribute), GILUsed(GILUsedAttribute), + #[cfg(feature = "experimental-module-state")] + State(StateAttribute), } impl Parse for PyModulePyO3Option { fn parse(input: ParseStream<'_>) -> Result { let lookahead = input.lookahead1(); + #[cfg(feature = "experimental-module-state")] + if lookahead.peek(attributes::kw::state) { + return input.parse().map(PyModulePyO3Option::State); + } + if lookahead.peek(attributes::kw::name) { input.parse().map(PyModulePyO3Option::Name) } else if lookahead.peek(syn::Token![crate]) { diff --git a/pyo3-macros/Cargo.toml b/pyo3-macros/Cargo.toml index 753117c7af2..559ad19b51b 100644 --- a/pyo3-macros/Cargo.toml +++ b/pyo3-macros/Cargo.toml @@ -18,6 +18,7 @@ proc-macro = true multiple-pymethods = [] experimental-async = ["pyo3-macros-backend/experimental-async"] experimental-inspect = ["pyo3-macros-backend/experimental-inspect"] +experimental-module-state = ["pyo3-macros-backend/experimental-module-state"] [dependencies] proc-macro2 = { version = "1.0.60", default-features = false } diff --git a/pyo3-macros/src/lib.rs b/pyo3-macros/src/lib.rs index bba07366b3c..cf5cfe59e6e 100644 --- a/pyo3-macros/src/lib.rs +++ b/pyo3-macros/src/lib.rs @@ -12,6 +12,62 @@ use pyo3_macros_backend::{ use quote::quote; use syn::{parse_macro_input, Item}; +/// Mark a struct as the module state container +/// +/// Used in `#[pymodule]` on modules to auto-detect the state type. +/// Requires the `experimental-module-state` feature to be enabled. +/// +/// # Example +/// ```ignore +/// #[pymodule] +/// mod my_module { +/// #[pymodule_state] +/// struct MyState { +/// cache: HashMap, +/// } +/// +/// fn init(m: &Bound<'_, PyModule>) -> PyResult { +/// Ok(MyState { cache: HashMap::new() }) +/// } +/// } +/// ``` +#[cfg(feature = "experimental-module-state")] +#[proc_macro_attribute] +pub fn pymodule_state( + _args: proc_macro::TokenStream, + input: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + // This macro does nothing - it's just a marker for #[pymodule]'s code generation + // The actual work is done by #[pymodule]'s analysis pass + input +} + +/// Mark a function as the module traverse handler for GC cycle detection. +/// Only one `#[pymodule_traverse]` allowed per module. +/// Requires: `experimental-module-state` feature. +#[cfg(feature = "experimental-module-state")] +#[proc_macro_attribute] +pub fn pymodule_traverse( + _args: proc_macro::TokenStream, + input: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + // This macro does nothing - it's just a marker for #[pymodule]'s code generation + input +} + +/// Mark a function as the module clear handler for breaking GC cycles. +/// Only one `#[pymodule_clear]` allowed per module. +/// Requires: `experimental-module-state` feature. +#[cfg(feature = "experimental-module-state")] +#[proc_macro_attribute] +pub fn pymodule_clear( + _args: proc_macro::TokenStream, + input: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + // This macro does nothing - it's just a marker for #[pymodule]'s code generation + input +} + /// A proc macro used to implement Python modules. /// /// The name of the module will be taken from the module name, unless `#[pyo3(name = "my_name")]` diff --git a/pytests/Cargo.toml b/pytests/Cargo.toml index 2f0d17102c2..277de4d7fe6 100644 --- a/pytests/Cargo.toml +++ b/pytests/Cargo.toml @@ -10,6 +10,7 @@ rust-version = "1.83" [features] experimental-async = ["pyo3/experimental-async"] experimental-inspect = ["pyo3/experimental-inspect"] +experimental-module-state = ["pyo3/experimental-module-state"] [dependencies] pyo3.path = "../" diff --git a/src/impl_.rs b/src/impl_.rs index 364f43ca4f8..9830f380e4c 100644 --- a/src/impl_.rs +++ b/src/impl_.rs @@ -28,3 +28,6 @@ pub mod pymodule; pub mod trampoline; pub mod unindent; pub mod wrap; + +#[cfg(feature = "experimental-module-state")] +pub mod pymodule_state; diff --git a/src/impl_/pyclass/lazy_type_object.rs b/src/impl_/pyclass/lazy_type_object.rs index 2142b1d4568..6dada0f2fe7 100644 --- a/src/impl_/pyclass/lazy_type_object.rs +++ b/src/impl_/pyclass/lazy_type_object.rs @@ -15,7 +15,7 @@ use crate::{ exceptions::PyRuntimeError, ffi, impl_::pymethods::PyMethodDefType, - pyclass::{create_type_object, PyClassTypeObject}, + pyclass::{create_type_object, create_type_object_with_module, PyClassTypeObject}, types::PyType, Bound, Py, PyAny, PyClass, PyErr, PyResult, Python, }; @@ -77,6 +77,62 @@ impl LazyTypeObject { T::items_iter(), ) } + + /// Gets the type object, initializing it with module association. + /// + /// # Errors + /// + /// Returns an error if: + /// - The type has already been created (race condition guard) + /// - Type creation fails + /// + /// This method enforces that the type is created with proper module context, + /// enabling `PyType_GetModule()` and `module_state()` to work correctly. + #[cold] + pub fn try_init_with_module<'py>( + &self, + py: Python<'py>, + module: *mut ffi::PyObject, + ) -> PyResult<&Bound<'py, PyType>> { + // Race condition guard: type must not already be created + if self.0.fully_initialized_type.get(py).is_some() { + return Err(PyErr::new::( + "Cannot create type with module: type already exists. \ + Call this method early in module initialization, \ + before the type is accessed elsewhere.", + )); + } + + // Manually perform initialization with module, similar to try_init but with module passed + (|| -> PyResult<_> { + let PyClassTypeObject { + type_object, + is_immutable_type, + .. + } = self + .0 + .value + .get_or_try_init(py, || create_type_object_with_module::(py, module))?; + let type_object = type_object.bind(py); + self.0.ensure_init( + type_object, + *is_immutable_type, + ::NAME, + T::items_iter(), + )?; + Ok(type_object) + })() + .map_err(|err| { + wrap_in_runtime_error( + py, + err, + format!( + "An error occurred while initializing class {}", + ::NAME + ), + ) + }) + } } impl LazyTypeObjectInner { diff --git a/src/impl_/pymodule.rs b/src/impl_/pymodule.rs index 03f7b356398..cccd15aa85f 100644 --- a/src/impl_/pymodule.rs +++ b/src/impl_/pymodule.rs @@ -4,6 +4,8 @@ //! Implementation details of `#[pymodule]` which need to be accessible from proc-macro generated code. #[allow(unused_imports, reason = "conditionally used")] use crate::platform::prelude::*; +#[cfg(feature = "experimental-module-state")] +use core::mem::MaybeUninit; use core::{ cell::UnsafeCell, ffi::CStr, @@ -32,19 +34,19 @@ use portable_atomic::AtomicI64; #[cfg(not(any(PyPy, GraalPy)))] use crate::exceptions::PyImportError; +#[cfg(feature = "experimental-module-state")] +use crate::exceptions::PyRuntimeError; +#[cfg(feature = "experimental-module-state")] +use crate::impl_::pymodule_state::ModuleState; use crate::prelude::PyTypeMethods; use crate::{ ffi, impl_::pyfunction::PyFunctionDef, - types::{PyModule, PyModuleMethods}, - Bound, PyClass, PyResult, PyTypeInfo, -}; -use crate::{ffi_ptr_ext::FfiPtrExt, PyErr}; -use crate::{ sync::PyOnceLock, - types::{any::PyAnyMethods, dict::PyDictMethods, PyDict}, - Py, PyAny, Python, + types::{any::PyAnyMethods, dict::PyDictMethods, PyDict, PyModule, PyModuleMethods}, + Bound, Py, PyAny, PyClass, PyResult, PyTypeInfo, Python, }; +use crate::{ffi_ptr_ext::FfiPtrExt, PyErr}; /// `Sync` wrapper of `ffi::PyModuleDef`. pub struct ModuleDef { @@ -75,7 +77,29 @@ impl ModuleDef { name: &'static CStr, doc: &'static CStr, slots: &'static PyModuleSlots, + allocate_state: bool, + traverse: Option, + clear: Option, ) -> Self { + #[cfg(feature = "experimental-module-state")] + let m_size = if allocate_state { + std::mem::size_of::() as _ + } else { + 0 + }; + #[cfg(feature = "experimental-module-state")] + let m_free = if allocate_state { + Some(pyo3_module_state_free as unsafe extern "C" fn(*mut c_void)) + } else { + None + }; + #[cfg(not(feature = "experimental-module-state"))] + let _ = allocate_state; + #[cfg(not(feature = "experimental-module-state"))] + let m_size = 0; + #[cfg(not(feature = "experimental-module-state"))] + let m_free = None; + // This is only used in PyO3 for append_to_inittab on Python 3.15 and newer. // There could also be other tools that need the legacy init hook. #[cfg(not(all(Py_LIMITED_API, Py_GIL_DISABLED)))] @@ -93,13 +117,19 @@ impl ModuleDef { }; #[cfg(not(all(Py_LIMITED_API, Py_GIL_DISABLED)))] - let ffi_def = UnsafeCell::new(ffi::PyModuleDef { - m_name: name.as_ptr(), - m_doc: doc.as_ptr(), - // TODO: would be slightly nicer to use `[T]::as_mut_ptr()` here, - // but that requires mut ptr deref on MSRV. - m_slots: slots.0.get() as _, - ..INIT + let ffi_def = UnsafeCell::new({ + ffi::PyModuleDef { + m_name: name.as_ptr(), + m_doc: doc.as_ptr(), + m_size, + // TODO: would be slightly nicer to use `[T]::as_mut_ptr()` here, + // but that requires mut ptr deref on MSRV. + m_slots: slots.0.get() as _, + m_free, + m_traverse: traverse, + m_clear: clear, + ..INIT + } }); ModuleDef { @@ -523,6 +553,38 @@ impl PyAddToModule for ModuleDef { } } +/// Called during multi-phase initialization in order to create an instance of +/// ModuleState on the memory area specific to modules. +#[cfg(feature = "experimental-module-state")] +pub fn pyo3_module_state_init(module: &Bound<'_, PyModule>, state: T) -> PyResult<()> { + unsafe { + let m_state: *mut MaybeUninit = ffi::PyModule_GetState(module.as_ptr()).cast(); + + // CPython builtins just assert this, but cross ffi panics are tricky, so we return an + // error instead + if m_state.is_null() { + return Err(PyRuntimeError::new_err( + "PyO3 per-module state was null. This is a bug in the Python interpreter runtime.", + )); + } + + (*m_state).write(ModuleState::new(state)); + + Ok(()) + } +} + +/// Called during deallocation of the module object. +/// +/// Used for the [`m_free`] field of [`PyModuleDef`]. +/// +/// [`m_free`]: https://docs.python.org/3/c-api/module.html#c.PyModuleDef.m_free +/// [`PyModuleDef`]: https://docs.python.org/3/c-api/module.html#c.PyModuleDef +#[cfg(feature = "experimental-module-state")] +pub unsafe extern "C" fn pyo3_module_state_free(module: *mut c_void) { + unsafe { ModuleState::pymodule_free_state(module.cast()) }; +} + #[cfg(test)] mod tests { use alloc::borrow::Cow; @@ -566,7 +628,7 @@ mod tests { .with_doc(DOC) .build(); - static MODULE_DEF: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS); + static MODULE_DEF: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS, false, None, None); Python::attach(|py| { let module = MODULE_DEF.make_module(py).unwrap().into_bound(py); @@ -606,7 +668,7 @@ mod tests { static SLOTS: PyModuleSlots = PyModuleSlotsBuilder::new().build(); - let module_def: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS); + let module_def: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS, false, None, None); #[cfg(not(all(Py_LIMITED_API, Py_GIL_DISABLED)))] unsafe { @@ -659,6 +721,128 @@ mod tests { assert!(result.is_err()); } + #[cfg(feature = "experimental-module-state")] + #[test] + fn module_state_init() { + use super::pyo3_module_state_init; + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + struct UserState(u64); + + static NAME: &CStr = c"test_module_state_init"; + static DOC: &CStr = c"This test is for checking PyO3 ModuleState is initialized correctly"; + + static SLOTS: PyModuleSlots = PyModuleSlotsBuilder::new().with_gil_used(false).build(); + + static MODULE_DEF: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS, true, None, None); + + Python::attach(|py| { + let module = MODULE_DEF + .make_module(py) + .expect("module to initialize without error") + .into_bound(py); + + let mystate = UserState(42); + + pyo3_module_state_init(&module, mystate).expect("state initialization should succeed"); + + assert_eq!( + Some(&mystate), + module.module_state::(), + "initialized state is referenceable" + ); + }) + } + + #[cfg(feature = "experimental-module-state")] + #[test] + fn module_state_type_mismatch() { + use super::pyo3_module_state_init; + + #[derive(Clone, Copy, PartialEq, Eq, Debug)] + struct StateA(i32); + + #[derive(Clone, Copy, PartialEq, Eq, Debug)] + struct StateB(i32); + + static NAME: &CStr = c"test_type_mismatch"; + static DOC: &CStr = c""; + static SLOTS: PyModuleSlots = PyModuleSlotsBuilder::new().build(); + static MODULE_DEF: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS, true, None, None); + + Python::attach(|py| { + let module = MODULE_DEF.make_module(py).unwrap().into_bound(py); + let state_a = StateA(42); + + pyo3_module_state_init(&module, state_a).unwrap(); + + // Type A accessible + assert_eq!(Some(&state_a), module.module_state::()); + + // Type B (different type) returns None - safe type mismatch handling + assert_eq!(None, module.module_state::()); + }) + } + + #[test] + fn module_state_without_allocation() { + static NAME: &CStr = c"test_no_state"; + static DOC: &CStr = c""; + static SLOTS: PyModuleSlots = PyModuleSlotsBuilder::new().build(); + + // allocate_state = false + static MODULE_DEF: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS, false, None, None); + + Python::attach(|py| { + let _module = MODULE_DEF.make_module(py).unwrap().into_bound(py); + + #[expect(dead_code)] + #[derive(Clone, Copy, PartialEq, Eq, Debug)] + struct AnyState(i32); + + // No state allocated, always returns None + #[cfg(feature = "experimental-module-state")] + assert_eq!(None, _module.module_state::()); + }) + } + + #[cfg(feature = "experimental-module-state")] + #[test] + fn module_state_mutable_access() { + use super::pyo3_module_state_init; + + #[derive(Clone, Copy, PartialEq, Eq, Debug)] + struct Counter { + value: i32, + } + + static NAME: &CStr = c"test_mutable"; + static DOC: &CStr = c""; + static SLOTS: PyModuleSlots = PyModuleSlotsBuilder::new().build(); + static MODULE_DEF: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS, true, None, None); + + Python::attach(|py| { + let mut module = MODULE_DEF.make_module(py).unwrap().into_bound(py); + + let initial = Counter { value: 0 }; + pyo3_module_state_init(&module, initial).unwrap(); + + // Verify initial + assert_eq!(Some(&initial), module.module_state::()); + + // Modify via mutable access + unsafe { + if let Some(state) = module.module_state_mut::() { + state.value = 42; + } + } + + // Verify modification persists + let expected = Counter { value: 42 }; + assert_eq!(Some(&expected), module.module_state::()); + }) + } + #[test] #[should_panic] fn test_module_slots_builder_overflow() { diff --git a/src/impl_/pymodule_state.rs b/src/impl_/pymodule_state.rs new file mode 100644 index 00000000000..b6a89a515ab --- /dev/null +++ b/src/impl_/pymodule_state.rs @@ -0,0 +1,134 @@ +use crate::types::PyModule; +use crate::{ffi, Bound}; +use alloc::boxed::Box; +use core::ptr::NonNull; + +/// Represents a Python module's state. +/// +/// More precisely, this `struct` resides on the per-module memory area +/// allocated during the module's creation. +#[repr(C)] +#[derive(Debug)] +pub struct ModuleState { + inner: Option>, +} + +impl ModuleState { + /// Create a new, empty [`ModuleState`] + pub fn new(value: T) -> Self { + Self { + inner: Some(Box::new(value)), + } + } + + /// Retrieve immutable reference to state + pub fn inner_ref(&self) -> Option<&T> { + self.inner.as_ref().and_then(|s| s.downcast_ref::()) + } + + /// Retrieve mutable reference to state + pub fn inner_mut(&mut self) -> Option<&mut T> { + self.inner.as_mut().and_then(|s| s.downcast_mut::()) + } + + /// This is the actual [`Drop::drop`] implementation, split out + /// so we can run it on the state ptr returned from [`Self::pymodule_get_state`] + /// + /// While this function does not take a owned `self`, the calling ModuleState + /// should not be accessed again + /// + /// Calling this function multiple times on a single ModuleState is a noop, + /// beyond the first + unsafe fn drop_impl(&mut self) { + self.inner.take(); + } +} + +impl ModuleState { + /// Fetch the [`ModuleState`] from a bound PyModule, inheriting it's lifetime + /// + /// ## Panics + /// + /// This function can panic if called on a PyModule that has not yet been + /// initialized + pub(crate) fn from_bound<'a>(this: &'a Bound<'_, PyModule>) -> Option<&'a Self> { + unsafe { Self::pymodule_get_state(this.as_ptr()).map(|ptr| ptr.as_ref()) } + } + + /// Fetch the [`ModuleState`] mutably from a bound PyModule, inheriting it's + /// lifetime + /// + /// ## Panics + /// + /// This function can panic if called on a PyModule that has not yet been + /// initialized + pub(crate) fn from_bound_mut<'a>(this: &'a mut Bound<'_, PyModule>) -> Option<&'a mut Self> { + unsafe { Self::pymodule_get_state(this.as_ptr()).map(|mut ptr| ptr.as_mut()) } + } + + /// Associated low level function for retrieving a pyo3 `pymodule`'s state + /// + /// If this function returns None, it means the underlying C PyModule does + /// not have module state. + /// + /// This function should only be called on a PyModule that is already + /// initialized via PyModule_New (or Py_mod_create) + pub(crate) unsafe fn pymodule_get_state(module: *mut ffi::PyObject) -> Option> { + unsafe { + let state: *mut ModuleState = ffi::PyModule_GetState(module).cast(); + + match state.is_null() { + true => None, + false => Some(NonNull::new_unchecked(state)), + } + } + } + + /// Associated low level function for freeing our `pymodule`'s state + /// via a ModuleDef's m_free C callback + pub(crate) unsafe fn pymodule_free_state(module: *mut ffi::PyObject) { + unsafe { + if let Some(state) = Self::pymodule_get_state(module) { + // SAFETY: this callback is called when python is freeing the + // associated PyModule, so we should never be accessed again + (*state.as_ptr()).drop_impl() + } + } + } + + /// Get the module state from a type's module. + /// + /// Returns `Ok(None)` if the module doesn't have state, state is not initialized, + /// or doesn't match the requested type. + /// + /// # Safety + /// + /// - `type_ptr` must be a valid, non-null pointer to a `PyTypeObject`. + /// - The GIL must be held (represented by `py`). + pub(crate) unsafe fn get_from_type_ptr<'py, T: 'static>( + py: crate::Python<'py>, + type_ptr: *mut ffi::PyTypeObject, + ) -> crate::PyResult> { + use crate::err; + + let mstate = unsafe { + let state_ptr = ffi::PyType_GetModuleState(type_ptr); + if state_ptr.is_null() { + if !ffi::PyErr_Occurred().is_null() { + return Err(err::PyErr::fetch(py)); + } + return Ok(None); + } + + &*(state_ptr as *const ModuleState) + }; + Ok(mstate.inner_ref::()) + } +} + +impl Drop for ModuleState { + fn drop(&mut self) { + // SAFETY: we're being dropped, so we'll never be accessed again + unsafe { self.drop_impl() }; + } +} diff --git a/src/marker.rs b/src/marker.rs index d764da8c651..fa44661bb97 100644 --- a/src/marker.rs +++ b/src/marker.rs @@ -654,6 +654,30 @@ impl<'py> Python<'py> { T::type_object(self) } + /// Gets the module state for a type. + /// + /// Returns `Ok(None)` if the type's module doesn't have state, the state is not initialized, + /// or doesn't match the requested state type. + /// + /// # Example + /// + /// ```ignore + /// if let Some(state) = py.get_module_state::()? { + /// // Use state + /// } + /// ``` + #[cfg(feature = "experimental-module-state")] + pub fn type_module_state(self) -> PyResult> + where + T: PyTypeInfo, + S: 'static, + { + use crate::impl_::pymodule_state::ModuleState; + let type_ptr = T::type_object_raw(self); + // Safety: `type_ptr` is a valid pointer to a Python type object + unsafe { ModuleState::get_from_type_ptr::(self, type_ptr) } + } + /// Imports the Python module with the specified name. pub fn import(self, name: N) -> PyResult> where diff --git a/src/pyclass.rs b/src/pyclass.rs index 10dbbec6bbf..92a3371abdc 100644 --- a/src/pyclass.rs +++ b/src/pyclass.rs @@ -6,7 +6,9 @@ mod create_type_object; mod gc; mod guard; -pub(crate) use self::create_type_object::{create_type_object, PyClassTypeObject}; +pub(crate) use self::create_type_object::{ + create_type_object, create_type_object_with_module, PyClassTypeObject, +}; pub use self::gc::{PyTraverseError, PyVisit}; pub use self::guard::{ diff --git a/src/pyclass/create_type_object.rs b/src/pyclass/create_type_object.rs index aa747262200..041cee0689b 100644 --- a/src/pyclass/create_type_object.rs +++ b/src/pyclass/create_type_object.rs @@ -44,6 +44,26 @@ pub(crate) struct PyClassTypeObject { } pub(crate) fn create_type_object(py: Python<'_>) -> PyResult +where + T: PyClass, +{ + create_type_object_inner::(py, None) +} + +pub(crate) fn create_type_object_with_module( + py: Python<'_>, + module: *mut ffi::PyObject, +) -> PyResult +where + T: PyClass, +{ + create_type_object_inner::(py, Some(module)) +} + +fn create_type_object_inner( + py: Python<'_>, + module_ptr: Option<*mut ffi::PyObject>, +) -> PyResult where T: PyClass, { @@ -67,6 +87,7 @@ where name: &'static str, module: Option<&'static str>, basicsize: ffi::Py_ssize_t, + module_ptr: Option<*mut ffi::PyObject>, ) -> PyResult { unsafe { PyTypeBuilder { @@ -97,7 +118,7 @@ where .offsets(dict_offset, weaklist_offset) .set_is_basetype(is_basetype) .class_items(items_iter) - .build(py, name, module, basicsize) + .build(py, name, module, basicsize, module_ptr) } } @@ -120,6 +141,7 @@ where ::NAME, ::MODULE, ::Layout::BASIC_SIZE, + module_ptr, ) } } @@ -414,6 +436,7 @@ impl PyTypeBuilder { name: &'static str, module_name: Option<&'static str>, basicsize: ffi::Py_ssize_t, + module_ptr: Option<*mut ffi::PyObject>, ) -> PyResult { // `c_ulong` and `c_uint` have the same size // on some platforms (like windows) @@ -503,27 +526,36 @@ impl PyTypeBuilder { #[cfg(Py_3_15)] let type_object = { - let mut slots = [ + // Use Vec to accommodate optional module slot in Python 3.15+ slot API + let mut slots_vec = vec![ ffi::PySlot_DATA(ffi::Py_tp_name, class_name.as_ptr() as *mut c_void), ffi::PySlot_UINT64(ffi::Py_tp_flags, flags.into()), ffi::PySlot_DATA(ffi::Py_tp_slots, self.slots.as_mut_ptr().cast::()), - match basicsize { - 1.. => ffi::PySlot_SIZE(ffi::Py_tp_basicsize, basicsize), - // zero size; don't set the slot at all, the VM will use the parent size - // - // This can *ONLY* be hit on the variable-sized case with a rust ZST, - // as a fully-sized case will always have size at least equal to `PyObject` - 0 => ffi::PySlot_END(), - ..0 => ffi::PySlot_SIZE(ffi::Py_tp_extra_basicsize, -basicsize), - }, - // NB: insert additional slots BEFORE `basicsize` slot as it might be null - ffi::PySlot_END(), ]; - // SAFETY: We've correctly setup the slots array at this point. - // The FFI call is known to return a new type object or null on error. + // Add basicsize slot if non-zero + match basicsize { + 1.. => slots_vec.push(ffi::PySlot_SIZE(ffi::Py_tp_basicsize, basicsize)), + // zero size; don't set the slot at all, the VM will use the parent size + // + // This can *ONLY* be hit on the variable-sized case with a rust ZST, + // as a fully-sized case will always have size at least equal to `PyObject` + 0 => {} + ..0 => slots_vec.push(ffi::PySlot_SIZE(ffi::Py_tp_extra_basicsize, -basicsize)), + } + + // In Python 3.15+, pass module context via Py_tp_module slot + if let Some(mod_ptr) = module_ptr { + slots_vec.push(ffi::PySlot_DATA(ffi::Py_tp_module, mod_ptr as *mut c_void)); + } + + // Sentinel to mark end of slots + slots_vec.push(ffi::PySlot_END()); + + // SAFETY: We've correctly setup the slots array. + // PyType_FromSlots returns a new type object or NULL on error. unsafe { - ffi::PyType_FromSlots(slots.as_mut_ptr()) + ffi::PyType_FromSlots(slots_vec.as_mut_ptr()) .assume_owned_or_err(py)? .cast_into_unchecked::() } @@ -542,9 +574,26 @@ impl PyTypeBuilder { // SAFETY: We've correctly setup the PyType_Spec at this point. // The FFI call is known to return a new type object or null on error. unsafe { - ffi::PyType_FromSpec(&mut spec) - .assume_owned_or_err(py)? - .cast_into_unchecked::() + // Use PyType_FromModuleAndSpec if module context available (Python 3.10+) + #[cfg(any(Py_3_10, not(Py_LIMITED_API)))] + if let Some(mod_ptr) = module_ptr { + ffi::PyType_FromModuleAndSpec(mod_ptr, &mut spec, ptr::null_mut()) + .assume_owned_or_err(py)? + .cast_into_unchecked::() + } else { + ffi::PyType_FromSpec(&mut spec) + .assume_owned_or_err(py)? + .cast_into_unchecked::() + } + + #[cfg(not(any(Py_3_10, not(Py_LIMITED_API))))] + { + let _ = module_ptr; + // Python < 3.10 with LIMITED_API: use fallback + ffi::PyType_FromSpec(&mut spec) + .assume_owned_or_err(py)? + .cast_into_unchecked::() + } } }; diff --git a/src/types/any.rs b/src/types/any.rs index 5821c0d4984..8e972a17bcd 100644 --- a/src/types/any.rs +++ b/src/types/any.rs @@ -820,6 +820,10 @@ pub trait PyAnyMethods<'py>: crate::sealed::Sealed { /// /// This is equivalent to the Python expression `super()` fn py_super(&self) -> PyResult>; + + // We do not provide a type_module_state mut since users cannot lock the module + // during access to the state, so it cannot be safe to access a mutable reference + // to the state. } macro_rules! implement_binop { diff --git a/src/types/module.rs b/src/types/module.rs index bea42e5d78e..0d1bebb0960 100644 --- a/src/types/module.rs +++ b/src/types/module.rs @@ -1,6 +1,8 @@ use crate::err::{PyErr, PyResult}; use crate::ffi_ptr_ext::FfiPtrExt; use crate::impl_::callback::IntoPyCallbackOutput; +#[cfg(feature = "experimental-module-state")] +use crate::impl_::pymodule_state::ModuleState; use crate::py_result_ext::PyResultExt; use crate::pyclass::PyClass; use crate::types::{ @@ -268,6 +270,16 @@ pub trait PyModuleMethods<'py>: crate::sealed::Sealed { /// Instead, this method is *generic*, and requires us to use the /// "turbofish" syntax to specify the class we want to add. /// + /// # Module Association and Module State + /// + /// When using module state, you need to use [`add_class_with_module`] instead. + /// If your class methods call `module_state()` or access `PyType_GetModuleState()`, + /// the type *must* be created with proper module context. This happens automatically + /// when using `add_class_with_module`, but `add_class` creates the type without + /// module context (if the type hasn't been accessed yet). + /// + /// For types that don't use module state, `add_class` is sufficient. + /// /// # Examples /// /// ```rust,no_run @@ -283,24 +295,40 @@ pub trait PyModuleMethods<'py>: crate::sealed::Sealed { /// } /// ``` /// - /// Python code can see this class as such: - /// ```python - /// from my_module import Foo + /// [`add_class_with_module`]: Self::add_class_with_module + #[doc = concat!("[1]: https://pyo3.rs/v", env!("CARGO_PKG_VERSION"), "/class.html#constructor")] + fn add_class(&self) -> PyResult<()> + where + T: PyClass; + + /// Adds a class to a module with module context. /// - /// print("Foo is", Foo) - /// ``` + /// This method should be used when you need `PyType_GetModuleState()` or `module_state()` + /// to work correctly in class methods. It ensures the type is created with the proper + /// module association, enabling module state access. /// - /// This will result in the following output: - /// ```text - /// Foo is - /// ``` + /// # Errors /// - /// Note that as we haven't defined a [constructor][1], Python code can't actually - /// make an *instance* of `Foo` (or *get* one for that matter, as we haven't exported - /// anything that can return instances of `Foo`). + /// Returns an error if the type has already been created (cached). To use this method, + /// you must call it *before* the class type is first accessed elsewhere in your code. + /// This is typically done early in the module initialization. /// - #[doc = concat!("[1]: https://pyo3.rs/v", env!("CARGO_PKG_VERSION"), "/class.html#constructor")] - fn add_class(&self) -> PyResult<()> + /// If you don't need module state in your class methods, prefer the regular [`add_class`] + /// method instead. + /// + /// # Example + /// + /// ```ignore + /// #[pymodule] + /// fn my_module(py: Python, m: &Bound) -> PyResult<()> { + /// // Add class with module context - enables module_state() calls + /// m.add_class_with_module::()?; + /// Ok(()) + /// } + /// ``` + /// + /// [`add_class`]: Self::add_class + fn add_class_with_module(&self) -> PyResult<()> where T: PyClass; @@ -419,6 +447,40 @@ pub trait PyModuleMethods<'py>: crate::sealed::Sealed { /// /// This is a no-op on the GIL-enabled build. fn gil_used(&self, gil_used: bool) -> PyResult<()>; + + /// Get a reference to the module state of type T + /// + /// Returns None if state is not initialized or type doesn't match. + /// + /// # Example + /// ```ignore + /// if let Some(state) = m.module_state::() { + /// println!("State: {:?}", state); + /// } + /// ``` + #[cfg(feature = "experimental-module-state")] + fn module_state(&self) -> Option<&T>; + + /// Get a mutable reference to the module state of type T + /// + /// Returns None if state is not initialized or type doesn't match. + /// + /// # Safety + /// + /// This is unsafe because it bypasses Rust's borrow checker. + /// You must ensure no other references exist to the state. + /// Locking the module in a critical section can be used to ensure this. + /// + /// # Example + /// ```ignore + /// unsafe { + /// if let Some(state) = m.module_state_mut::() { + /// state.initialize()?; + /// } + /// } + /// ``` + #[cfg(feature = "experimental-module-state")] + unsafe fn module_state_mut(&mut self) -> Option<&mut T>; } impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> { @@ -522,6 +584,17 @@ impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> { ) } + fn add_class_with_module(&self) -> PyResult<()> + where + T: PyClass, + { + let py = self.py(); + self.add( + ::NAME, + T::lazy_type_object().try_init_with_module(py, self.as_ptr())?, + ) + } + fn add_wrapped(&self, wrapper: &impl Fn(Python<'py>) -> T) -> PyResult<()> where T: IntoPyCallbackOutput<'py, Py>, @@ -566,6 +639,16 @@ impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> { #[cfg(any(Py_LIMITED_API, not(Py_GIL_DISABLED)))] Ok(()) } + + #[cfg(feature = "experimental-module-state")] + fn module_state(&self) -> Option<&T> { + ModuleState::from_bound(self).and_then(|state| state.inner_ref::()) + } + + #[cfg(feature = "experimental-module-state")] + unsafe fn module_state_mut(&mut self) -> Option<&mut T> { + ModuleState::from_bound_mut(self).and_then(|state| state.inner_mut::()) + } } fn __all__(py: Python<'_>) -> &Bound<'_, PyString> { diff --git a/src/types/typeobject.rs b/src/types/typeobject.rs index 5003ee0205a..71e8c65aa85 100644 --- a/src/types/typeobject.rs +++ b/src/types/typeobject.rs @@ -4,6 +4,8 @@ use crate::instance::Borrowed; use crate::pybacked::PyBackedStr; #[cfg(any(Py_LIMITED_API, PyPy, not(Py_3_13)))] use crate::types::any::PyAnyMethods; +#[cfg(any(not(Py_LIMITED_API), Py_3_10))] +use crate::types::PyModule; use crate::types::PyTuple; use crate::{ffi, Bound, PyAny, PyTypeInfo, Python}; #[cfg(RustPython)] @@ -108,6 +110,23 @@ pub trait PyTypeMethods<'py>: crate::sealed::Sealed { /// /// Equivalent to the Python expression `self.__bases__`. fn bases(&self) -> Bound<'py, PyTuple>; + + /// Get the module object that defines this type. + /// + /// # Errors + /// Returns an error if the type doesn't have a module. + #[cfg(any(not(Py_LIMITED_API), Py_3_10))] + fn module_object(&self) -> PyResult>; + + /// Get the module state from this type's module. + /// + /// Returns `Ok(None)` if the module doesn't have state. + /// + /// # Errors + /// Returns an error if the type doesn't have a module or if state retrieval fails. + #[cfg(feature = "experimental-module-state")] + #[cfg(any(not(Py_LIMITED_API), Py_3_10))] + fn module_state(&self) -> PyResult>; } impl<'py> PyTypeMethods<'py> for Bound<'py, PyType> { @@ -255,6 +274,25 @@ impl<'py> PyTypeMethods<'py> for Bound<'py, PyType> { bases } + + #[cfg(any(not(Py_LIMITED_API), Py_3_10))] + fn module_object(&self) -> PyResult> { + unsafe { + let mod_ptr = ffi::PyType_GetModule(self.as_type_ptr()); + if mod_ptr.is_null() { + return Err(err::PyErr::fetch(self.py())); + } + Ok(Bound::from_borrowed_ptr(self.py(), mod_ptr).cast_into_unchecked::()) + } + } + + #[cfg(feature = "experimental-module-state")] + #[cfg(any(not(Py_LIMITED_API), Py_3_10))] + fn module_state(&self) -> PyResult> { + use crate::impl_::pymodule_state::ModuleState; + // SAFETY: `self.as_type_ptr()` is a valid pointer to a Python type object + unsafe { ModuleState::get_from_type_ptr::(self.py(), self.as_type_ptr()) } + } } #[cfg(test)]