Skip to content

Commit fee1a4b

Browse files
youknowoneyouknowone
committed
multiphase
1 parent 3ccff6d commit fee1a4b

File tree

8 files changed

+141
-13
lines changed

8 files changed

+141
-13
lines changed

‎crates/stdlib/src/lib.rs‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,9 +105,23 @@ mod tkinter;
105105
use rustpython_common as common;
106106
use rustpython_vm as vm;
107107

108-
use crate::vm::{builtins, stdlib::StdlibInitFunc};
108+
use crate::vm::{
109+
builtins,
110+
stdlib::{StdlibDefFunc, StdlibInitFunc},
111+
};
109112
use alloc::borrow::Cow;
110113

114+
/// Returns module definitions for multi-phase init modules.
115+
/// These modules are added to sys.modules BEFORE their exec function runs,
116+
/// allowing safe circular imports.
117+
pub fn get_module_defs() -> impl Iterator<Item = (Cow<'static, str>, StdlibDefFunc)> {
118+
[(
119+
Cow::from("_asyncio"),
120+
_asyncio::__module_def as StdlibDefFunc,
121+
)]
122+
.into_iter()
123+
}
124+
111125
pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInitFunc)> {
112126
macro_rules! modules {
113127
{

‎crates/vm/src/builtins/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ pub use mappingproxy::PyMappingProxy;
4949
pub(crate) mod memory;
5050
pub use memory::PyMemoryView;
5151
pub(crate) mod module;
52-
pub use module::{PyModule, PyModuleDef};
52+
pub use module::{PyModule, PyModuleDef, PyModuleSlots};
5353
pub(crate) mod namespace;
5454
pub use namespace::PyNamespace;
5555
pub(crate) mod object;

‎crates/vm/src/import.rs‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,14 +87,50 @@ pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
8787
}
8888

8989
pub fn import_builtin(vm: &VirtualMachine, module_name: &str) -> PyResult {
90+
use crate::PyPayload;
91+
use crate::builtins::PyModule;
92+
93+
let sys_modules = vm.sys_module.get_attr("modules", vm)?;
94+
95+
// Check if already in sys.modules (handles recursive imports)
96+
if let Ok(module) = sys_modules.get_item(module_name, vm) {
97+
return Ok(module);
98+
}
99+
100+
// Try multi-phase init first (preferred for modules that import other modules)
101+
if let Some(def_func) = vm.state.module_defs.get(module_name) {
102+
let def = def_func(&vm.ctx);
103+
104+
// Phase 1: Create module from definition
105+
let module = PyModule::from_def(def).into_ref(&vm.ctx);
106+
107+
// Initialize module dict
108+
let dict = vm.ctx.new_dict();
109+
dict.set_item("__name__", vm.ctx.new_str(def.name.as_str()).into(), vm)?;
110+
if let Some(doc) = def.doc {
111+
dict.set_item("__doc__", vm.ctx.new_str(doc.as_str()).into(), vm)?;
112+
}
113+
module.set_attr("__dict__", dict, vm)?;
114+
115+
// Add to sys.modules BEFORE exec (critical for circular import handling)
116+
sys_modules.set_item(module_name, module.clone().into(), vm)?;
117+
118+
// Phase 2: Call exec slot (can safely import other modules now)
119+
if let Some(exec) = def.slots.exec {
120+
exec(vm, &module)?;
121+
}
122+
123+
return Ok(module.into());
124+
}
125+
126+
// Fall back to legacy single-phase init
90127
let make_module_func = vm.state.module_inits.get(module_name).ok_or_else(|| {
91128
vm.new_import_error(
92129
format!("Cannot import builtin module {module_name}"),
93130
vm.ctx.new_str(module_name),
94131
)
95132
})?;
96133
let module = make_module_func(vm);
97-
let sys_modules = vm.sys_module.get_attr("modules", vm)?;
98134
sys_modules.set_item(module_name, module.as_object().to_owned(), vm)?;
99135
Ok(module.into())
100136
}

‎crates/vm/src/stdlib/imp.rs‎

Lines changed: 44 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ mod _imp {
107107
#[pyfunction]
108108
fn is_builtin(name: PyStrRef, vm: &VirtualMachine) -> bool {
109109
vm.state.module_inits.contains_key(name.as_str())
110+
|| vm.state.module_defs.contains_key(name.as_str())
110111
}
111112

112113
#[pyfunction]
@@ -116,22 +117,56 @@ mod _imp {
116117

117118
#[pyfunction]
118119
fn create_builtin(spec: PyObjectRef, vm: &VirtualMachine) -> PyResult {
120+
use crate::PyPayload;
121+
use crate::builtins::PyModule;
122+
119123
let sys_modules = vm.sys_module.get_attr("modules", vm).unwrap();
120124
let name: PyStrRef = spec.get_attr("name", vm)?.try_into_value(vm)?;
121125

122-
let module = if let Ok(module) = sys_modules.get_item(&*name, vm) {
123-
module
124-
} else if let Some(make_module_func) = vm.state.module_inits.get(name.as_str()) {
125-
make_module_func(vm).into()
126-
} else {
127-
vm.ctx.none()
128-
};
129-
Ok(module)
126+
// Check sys.modules first
127+
if let Ok(module) = sys_modules.get_item(&*name, vm) {
128+
return Ok(module);
129+
}
130+
131+
// Try multi-phase init modules first (they need special handling)
132+
if let Some(def_func) = vm.state.module_defs.get(name.as_str()) {
133+
let def = def_func(&vm.ctx);
134+
135+
// Phase 1: Create module from definition
136+
let module = PyModule::from_def(def).into_ref(&vm.ctx);
137+
138+
// Initialize module dict
139+
let dict = vm.ctx.new_dict();
140+
dict.set_item("__name__", vm.ctx.new_str(def.name.as_str()).into(), vm)?;
141+
if let Some(doc) = def.doc {
142+
dict.set_item("__doc__", vm.ctx.new_str(doc.as_str()).into(), vm)?;
143+
}
144+
module.set_attr("__dict__", dict, vm)?;
145+
146+
// Add to sys.modules BEFORE exec (critical for circular import handling)
147+
sys_modules.set_item(&*name, module.clone().into(), vm)?;
148+
149+
// Phase 2: Call exec slot (can safely import other modules now)
150+
if let Some(exec) = def.slots.exec {
151+
exec(vm, &module)?;
152+
}
153+
154+
return Ok(module.into());
155+
}
156+
157+
// Fall back to legacy single-phase init
158+
if let Some(make_module_func) = vm.state.module_inits.get(name.as_str()) {
159+
let module = make_module_func(vm);
160+
sys_modules.set_item(&*name, module.clone().into(), vm)?;
161+
return Ok(module.into());
162+
}
163+
164+
Ok(vm.ctx.none())
130165
}
131166

132167
#[pyfunction]
133168
fn exec_builtin(_mod: PyRef<PyModule>) -> i32 {
134-
// TODO: Should we do something here?
169+
// For multi-phase init modules, exec is already called in create_builtin
135170
0
136171
}
137172

‎crates/vm/src/stdlib/mod.rs‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,12 +62,22 @@ mod winapi;
6262
#[cfg(windows)]
6363
mod winreg;
6464

65-
use crate::{PyRef, VirtualMachine, builtins::PyModule};
65+
use crate::{Context, PyRef, VirtualMachine, builtins::PyModule, builtins::PyModuleDef};
6666
use alloc::borrow::Cow;
6767
use std::collections::HashMap;
6868

69+
/// Legacy single-phase init: function that creates and populates a module in one step
6970
pub type StdlibInitFunc = Box<py_dyn_fn!(dyn Fn(&VirtualMachine) -> PyRef<PyModule>)>;
71+
72+
/// Multi-phase init: function that returns a module definition
73+
/// The import machinery will:
74+
/// 1. Create module from def
75+
/// 2. Add to sys.modules
76+
/// 3. Call exec slot (which can safely import other modules)
77+
pub type StdlibDefFunc = fn(&Context) -> &'static PyModuleDef;
78+
7079
pub type StdlibMap = HashMap<Cow<'static, str>, StdlibInitFunc, ahash::RandomState>;
80+
pub type StdlibDefMap = HashMap<Cow<'static, str>, StdlibDefFunc, ahash::RandomState>;
7181

7282
pub fn get_module_inits() -> StdlibMap {
7383
macro_rules! modules {
@@ -154,3 +164,12 @@ pub fn get_module_inits() -> StdlibMap {
154164
}
155165
}
156166
}
167+
168+
/// Returns module definitions for multi-phase init modules.
169+
/// These modules use CPython's two-phase initialization pattern:
170+
/// 1. Create module from def and add to sys.modules
171+
/// 2. Call exec slot (can safely import other modules without circular import issues)
172+
pub fn get_module_defs() -> StdlibDefMap {
173+
// Currently empty - modules will be migrated to multi-phase init as needed
174+
HashMap::default()
175+
}

‎crates/vm/src/stdlib/sys.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,8 @@ mod sys {
169169
#[pyattr]
170170
fn builtin_module_names(vm: &VirtualMachine) -> PyTupleRef {
171171
let mut module_names: Vec<_> = vm.state.module_inits.keys().cloned().collect();
172+
// Also include multi-phase init modules
173+
module_names.extend(vm.state.module_defs.keys().cloned());
172174
module_names.push("sys".into());
173175
module_names.push("builtins".into());
174176
module_names.sort();

‎crates/vm/src/vm/mod.rs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ struct ExceptionStack {
103103
pub struct PyGlobalState {
104104
pub config: PyConfig,
105105
pub module_inits: stdlib::StdlibMap,
106+
pub module_defs: stdlib::StdlibDefMap,
106107
pub frozen: HashMap<&'static str, FrozenModule, ahash::RandomState>,
107108
pub stacksize: AtomicCell<usize>,
108109
pub thread_count: AtomicCell<usize>,
@@ -170,6 +171,7 @@ impl VirtualMachine {
170171
));
171172

172173
let module_inits = stdlib::get_module_inits();
174+
let module_defs = stdlib::get_module_defs();
173175

174176
let seed = match config.settings.hash_seed {
175177
Some(seed) => seed,
@@ -203,6 +205,7 @@ impl VirtualMachine {
203205
state: PyRc::new(PyGlobalState {
204206
config,
205207
module_inits,
208+
module_defs,
206209
frozen: HashMap::default(),
207210
stacksize: AtomicCell::new(0),
208211
thread_count: AtomicCell::new(0),
@@ -490,6 +493,23 @@ impl VirtualMachine {
490493
self.state_mut().module_inits.extend(iter);
491494
}
492495

496+
/// Add a module definition for multi-phase initialization.
497+
/// These modules are added to sys.modules BEFORE their exec function runs,
498+
/// allowing safe circular imports.
499+
pub fn add_native_module_def<S>(&mut self, name: S, def_func: stdlib::StdlibDefFunc)
500+
where
501+
S: Into<Cow<'static, str>>,
502+
{
503+
self.state_mut().module_defs.insert(name.into(), def_func);
504+
}
505+
506+
pub fn add_native_module_defs<I>(&mut self, iter: I)
507+
where
508+
I: IntoIterator<Item = (Cow<'static, str>, stdlib::StdlibDefFunc)>,
509+
{
510+
self.state_mut().module_defs.extend(iter);
511+
}
512+
493513
/// Can only be used in the initialization closure passed to [`Interpreter::with_init`]
494514
pub fn add_frozen<I>(&mut self, frozen: I)
495515
where

‎src/interpreter.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,8 @@ impl InterpreterConfig {
101101
/// Initializes all standard library modules for the given VM.
102102
#[cfg(feature = "stdlib")]
103103
pub fn init_stdlib(vm: &mut VirtualMachine) {
104+
// Add multi-phase init modules first (they're checked first in import_builtin)
105+
vm.add_native_module_defs(rustpython_stdlib::get_module_defs());
104106
vm.add_native_modules(rustpython_stdlib::get_module_inits());
105107

106108
#[cfg(feature = "freeze-stdlib")]

0 commit comments

Comments
 (0)