diff --git a/src/lib.rs b/src/lib.rs index ec8e11f18..9aca42ee0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -304,12 +304,7 @@ fn parse_arguments<'a>(app: App<'a, '_>) -> ArgMatches<'a> { fn add_stdlib(vm: &mut VirtualMachine) { let _ = vm; #[cfg(feature = "stdlib")] - { - let stdlib = rustpython_stdlib::get_module_inits(); - for (name, init) in stdlib.into_iter() { - vm.add_native_module(name, init); - } - } + vm.add_native_modules(rustpython_stdlib::get_module_inits()); } /// Create settings by examining command line arguments and environment diff --git a/stdlib/src/lib.rs b/stdlib/src/lib.rs index cff9fce54..0b9ea553f 100644 --- a/stdlib/src/lib.rs +++ b/stdlib/src/lib.rs @@ -52,13 +52,10 @@ mod termios; use rustpython_common as common; use rustpython_vm as vm; -use crate::vm::{ - builtins, - stdlib::{StdlibInitFunc, StdlibMap}, -}; +use crate::vm::{builtins, stdlib::StdlibInitFunc}; use std::borrow::Cow; -pub fn get_module_inits() -> StdlibMap { +pub fn get_module_inits() -> impl Iterator, StdlibInitFunc)> { macro_rules! modules { { $( @@ -66,12 +63,12 @@ pub fn get_module_inits() -> StdlibMap { { $( $key:expr => $val:expr),* $(,)? } )* } => {{ - let modules = [ + [ $( $(#[cfg($cfg)] (Cow::<'static, str>::from($key), Box::new($val) as StdlibInitFunc),)* )* - ]; - modules.into_iter().collect() + ] + .into_iter() }}; } modules! { diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 017391dd1..04a98a9cc 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -386,14 +386,24 @@ impl VirtualMachine { self.initialized = true; } + fn state_mut(&mut self) -> &mut PyGlobalState { + PyRc::get_mut(&mut self.state) + .expect("there should not be multiple threads while a user has a mut ref to a vm") + } + /// Can only be used in the initialization closure passed to [`Interpreter::new_with_init`] pub fn add_native_module(&mut self, name: S, module: stdlib::StdlibInitFunc) where S: Into>, { - let state = PyRc::get_mut(&mut self.state) - .expect("can't add_native_module when there are multiple threads"); - state.module_inits.insert(name.into(), module); + self.state_mut().module_inits.insert(name.into(), module); + } + + pub fn add_native_modules(&mut self, iter: I) + where + I: IntoIterator, stdlib::StdlibInitFunc)>, + { + self.state_mut().module_inits.extend(iter); } /// Can only be used in the initialization closure passed to [`Interpreter::new_with_init`] @@ -402,9 +412,7 @@ impl VirtualMachine { I: IntoIterator, { let frozen = frozen::map_frozen(self, frozen).collect::>(); - let state = PyRc::get_mut(&mut self.state) - .expect("can't add_frozen when there are multiple threads"); - state.frozen.extend(frozen); + self.state_mut().frozen.extend(frozen); } /// Start a new thread with access to the same interpreter.