rustpython_vm/vm/
mod.rs

1//! Implement virtual machine to run instructions.
2//!
3//! See also:
4//!   <https://github.com/ProgVal/pythonvm-rust/blob/master/src/processor/mod.rs>
5
6#[cfg(feature = "rustpython-compiler")]
7mod compile;
8mod context;
9mod interpreter;
10mod method;
11mod setting;
12pub mod thread;
13mod vm_new;
14mod vm_object;
15mod vm_ops;
16
17use crate::{
18    builtins::{
19        code::PyCode,
20        pystr::AsPyStr,
21        tuple::{PyTuple, PyTupleTyped},
22        PyBaseExceptionRef, PyDictRef, PyInt, PyList, PyModule, PyStr, PyStrInterned, PyStrRef,
23        PyTypeRef,
24    },
25    codecs::CodecsRegistry,
26    common::{hash::HashSecret, lock::PyMutex, rc::PyRc},
27    convert::ToPyObject,
28    frame::{ExecutionResult, Frame, FrameRef},
29    frozen::FrozenModule,
30    function::{ArgMapping, FuncArgs, PySetterValue},
31    import,
32    protocol::PyIterIter,
33    scope::Scope,
34    signal, stdlib,
35    warn::WarningsState,
36    AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
37};
38use crossbeam_utils::atomic::AtomicCell;
39#[cfg(unix)]
40use nix::{
41    sys::signal::{kill, sigaction, SaFlags, SigAction, SigSet, Signal::SIGINT},
42    unistd::getpid,
43};
44use std::sync::atomic::AtomicBool;
45use std::{
46    borrow::Cow,
47    cell::{Cell, Ref, RefCell},
48    collections::{HashMap, HashSet},
49};
50
51pub use context::Context;
52pub use interpreter::Interpreter;
53pub(crate) use method::PyMethod;
54pub use setting::Settings;
55
56// Objects are live when they are on stack, or referenced by a name (for now)
57
58/// Top level container of a python virtual machine. In theory you could
59/// create more instances of this struct and have them operate fully isolated.
60///
61/// To construct this, please refer to the [`Interpreter`](Interpreter)
62pub struct VirtualMachine {
63    pub builtins: PyRef<PyModule>,
64    pub sys_module: PyRef<PyModule>,
65    pub ctx: PyRc<Context>,
66    pub frames: RefCell<Vec<FrameRef>>,
67    pub wasm_id: Option<String>,
68    exceptions: RefCell<ExceptionStack>,
69    pub import_func: PyObjectRef,
70    pub profile_func: RefCell<PyObjectRef>,
71    pub trace_func: RefCell<PyObjectRef>,
72    pub use_tracing: Cell<bool>,
73    pub recursion_limit: Cell<usize>,
74    pub(crate) signal_handlers: Option<Box<RefCell<[Option<PyObjectRef>; signal::NSIG]>>>,
75    pub(crate) signal_rx: Option<signal::UserSignalReceiver>,
76    pub repr_guards: RefCell<HashSet<usize>>,
77    pub state: PyRc<PyGlobalState>,
78    pub initialized: bool,
79    recursion_depth: Cell<usize>,
80}
81
82#[derive(Debug, Default)]
83struct ExceptionStack {
84    exc: Option<PyBaseExceptionRef>,
85    prev: Option<Box<ExceptionStack>>,
86}
87
88pub struct PyGlobalState {
89    pub settings: Settings,
90    pub module_inits: stdlib::StdlibMap,
91    pub frozen: HashMap<&'static str, FrozenModule, ahash::RandomState>,
92    pub stacksize: AtomicCell<usize>,
93    pub thread_count: AtomicCell<usize>,
94    pub hash_secret: HashSecret,
95    pub atexit_funcs: PyMutex<Vec<(PyObjectRef, FuncArgs)>>,
96    pub codec_registry: CodecsRegistry,
97    pub finalizing: AtomicBool,
98    pub warnings: WarningsState,
99    pub override_frozen_modules: AtomicCell<isize>,
100    pub before_forkers: PyMutex<Vec<PyObjectRef>>,
101    pub after_forkers_child: PyMutex<Vec<PyObjectRef>>,
102    pub after_forkers_parent: PyMutex<Vec<PyObjectRef>>,
103    pub int_max_str_digits: AtomicCell<usize>,
104}
105
106pub fn process_hash_secret_seed() -> u32 {
107    use once_cell::sync::OnceCell;
108    static SEED: OnceCell<u32> = OnceCell::new();
109    *SEED.get_or_init(rand::random)
110}
111
112impl VirtualMachine {
113    /// Create a new `VirtualMachine` structure.
114    fn new(settings: Settings, ctx: PyRc<Context>) -> VirtualMachine {
115        flame_guard!("new VirtualMachine");
116
117        // make a new module without access to the vm; doesn't
118        // set __spec__, __loader__, etc. attributes
119        let new_module = |def| {
120            PyRef::new_ref(
121                PyModule::from_def(def),
122                ctx.types.module_type.to_owned(),
123                Some(ctx.new_dict()),
124            )
125        };
126
127        // Hard-core modules:
128        let builtins = new_module(stdlib::builtins::__module_def(&ctx));
129        let sys_module = new_module(stdlib::sys::__module_def(&ctx));
130
131        let import_func = ctx.none();
132        let profile_func = RefCell::new(ctx.none());
133        let trace_func = RefCell::new(ctx.none());
134        // hack to get around const array repeat expressions, rust issue #79270
135        const NONE: Option<PyObjectRef> = None;
136        // putting it in a const optimizes better, prevents linear initialization of the array
137        #[allow(clippy::declare_interior_mutable_const)]
138        const SIGNAL_HANDLERS: RefCell<[Option<PyObjectRef>; signal::NSIG]> =
139            RefCell::new([NONE; signal::NSIG]);
140        let signal_handlers = Some(Box::new(SIGNAL_HANDLERS));
141
142        let module_inits = stdlib::get_module_inits();
143
144        let seed = match settings.hash_seed {
145            Some(seed) => seed,
146            None => process_hash_secret_seed(),
147        };
148        let hash_secret = HashSecret::new(seed);
149
150        let codec_registry = CodecsRegistry::new(&ctx);
151
152        let warnings = WarningsState::init_state(&ctx);
153
154        let int_max_str_digits = AtomicCell::new(match settings.int_max_str_digits {
155            -1 => 4300,
156            other => other,
157        } as usize);
158        let mut vm = VirtualMachine {
159            builtins,
160            sys_module,
161            ctx,
162            frames: RefCell::new(vec![]),
163            wasm_id: None,
164            exceptions: RefCell::default(),
165            import_func,
166            profile_func,
167            trace_func,
168            use_tracing: Cell::new(false),
169            recursion_limit: Cell::new(if cfg!(debug_assertions) { 256 } else { 1000 }),
170            signal_handlers,
171            signal_rx: None,
172            repr_guards: RefCell::default(),
173            state: PyRc::new(PyGlobalState {
174                settings,
175                module_inits,
176                frozen: HashMap::default(),
177                stacksize: AtomicCell::new(0),
178                thread_count: AtomicCell::new(0),
179                hash_secret,
180                atexit_funcs: PyMutex::default(),
181                codec_registry,
182                finalizing: AtomicBool::new(false),
183                warnings,
184                override_frozen_modules: AtomicCell::new(0),
185                before_forkers: PyMutex::default(),
186                after_forkers_child: PyMutex::default(),
187                after_forkers_parent: PyMutex::default(),
188                int_max_str_digits,
189            }),
190            initialized: false,
191            recursion_depth: Cell::new(0),
192        };
193
194        if vm.state.hash_secret.hash_str("")
195            != vm
196                .ctx
197                .interned_str("")
198                .expect("empty str must be interned")
199                .hash(&vm)
200        {
201            panic!("Interpreters in same process must share the hash seed");
202        }
203
204        let frozen = core_frozen_inits().collect();
205        PyRc::get_mut(&mut vm.state).unwrap().frozen = frozen;
206
207        vm.builtins.init_dict(
208            vm.ctx.intern_str("builtins"),
209            Some(vm.ctx.intern_str(stdlib::builtins::DOC.unwrap()).to_owned()),
210            &vm,
211        );
212        vm.sys_module.init_dict(
213            vm.ctx.intern_str("sys"),
214            Some(vm.ctx.intern_str(stdlib::sys::DOC.unwrap()).to_owned()),
215            &vm,
216        );
217        // let name = vm.sys_module.get_attr("__name__", &vm).unwrap();
218        vm
219    }
220
221    /// set up the encodings search function
222    /// init_importlib must be called before this call
223    #[cfg(feature = "encodings")]
224    fn import_encodings(&mut self) -> PyResult<()> {
225        self.import("encodings", 0).map_err(|import_err| {
226            let rustpythonpath_env = std::env::var("RUSTPYTHONPATH").ok();
227            let pythonpath_env = std::env::var("PYTHONPATH").ok();
228            let env_set = rustpythonpath_env.as_ref().is_some() || pythonpath_env.as_ref().is_some();
229            let path_contains_env = self.state.settings.path_list.iter().any(|s| {
230                Some(s.as_str()) == rustpythonpath_env.as_deref() || Some(s.as_str()) == pythonpath_env.as_deref()
231            });
232
233            let guide_message = if cfg!(feature = "freeze-stdlib") {
234                "`rustpython_pylib` maybe not set while using `freeze-stdlib` feature. Try using `rustpython::InterpreterConfig::init_stdlib` or manually call `vm.add_frozen(rustpython_pylib::FROZEN_STDLIB)` in `rustpython_vm::Interpreter::with_init`."
235            } else if !env_set {
236                "Neither RUSTPYTHONPATH nor PYTHONPATH is set. Try setting one of them to the stdlib directory."
237            } else if path_contains_env {
238                "RUSTPYTHONPATH or PYTHONPATH is set, but it doesn't contain the encodings library. If you are customizing the RustPython vm/interpreter, try adding the stdlib directory to the path. If you are developing the RustPython interpreter, it might be a bug during development."
239            } else {
240                "RUSTPYTHONPATH or PYTHONPATH is set, but it wasn't loaded to `Settings::path_list`. If you are going to customize the RustPython vm/interpreter, those environment variables are not loaded in the Settings struct by default. Please try creating a customized instance of the Settings struct. If you are developing the RustPython interpreter, it might be a bug during development."
241            };
242
243            let mut msg = format!(
244                "RustPython could not import the encodings module. It usually means something went wrong. Please carefully read the following messages and follow the steps.\n\
245                \n\
246                {guide_message}");
247            if !cfg!(feature = "freeze-stdlib") {
248                msg += "\n\
249                If you don't have access to a consistent external environment (e.g. targeting wasm, embedding \
250                    rustpython in another application), try enabling the `freeze-stdlib` feature.\n\
251                If this is intended and you want to exclude the encodings module from your interpreter, please remove the `encodings` feature from `rustpython-vm` crate.";
252            }
253
254            let err = self.new_runtime_error(msg);
255            err.set_cause(Some(import_err));
256            err
257        })?;
258        Ok(())
259    }
260
261    fn import_utf8_encodings(&mut self) -> PyResult<()> {
262        import::import_frozen(self, "codecs")?;
263        // FIXME: See corresponding part of `core_frozen_inits`
264        // let encoding_module_name = if cfg!(feature = "freeze-stdlib") {
265        //     "encodings.utf_8"
266        // } else {
267        //     "encodings_utf_8"
268        // };
269        let encoding_module_name = "encodings_utf_8";
270        let encoding_module = import::import_frozen(self, encoding_module_name)?;
271        let getregentry = encoding_module.get_attr("getregentry", self)?;
272        let codec_info = getregentry.call((), self)?;
273        self.state
274            .codec_registry
275            .register_manual("utf-8", codec_info.try_into_value(self)?)?;
276        Ok(())
277    }
278
279    fn initialize(&mut self) {
280        flame_guard!("init VirtualMachine");
281
282        if self.initialized {
283            panic!("Double Initialize Error");
284        }
285
286        stdlib::builtins::init_module(self, &self.builtins);
287        stdlib::sys::init_module(self, &self.sys_module, &self.builtins);
288
289        let mut essential_init = || -> PyResult {
290            import::import_builtin(self, "_typing")?;
291            #[cfg(not(target_arch = "wasm32"))]
292            import::import_builtin(self, "_signal")?;
293            #[cfg(any(feature = "parser", feature = "compiler"))]
294            import::import_builtin(self, "_ast")?;
295            #[cfg(not(feature = "threading"))]
296            import::import_frozen(self, "_thread")?;
297            let importlib = import::init_importlib_base(self)?;
298            self.import_utf8_encodings()?;
299
300            #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))]
301            {
302                // this isn't fully compatible with CPython; it imports "io" and sets
303                // builtins.open to io.OpenWrapper, but this is easier, since it doesn't
304                // require the Python stdlib to be present
305                let io = import::import_builtin(self, "_io")?;
306                let set_stdio = |name, fd, mode: &str| {
307                    let stdio = crate::stdlib::io::open(
308                        self.ctx.new_int(fd).into(),
309                        Some(mode),
310                        Default::default(),
311                        self,
312                    )?;
313                    let dunder_name = self.ctx.intern_str(format!("__{name}__"));
314                    self.sys_module.set_attr(
315                        dunder_name, // e.g. __stdin__
316                        stdio.clone(),
317                        self,
318                    )?;
319                    self.sys_module.set_attr(name, stdio, self)?;
320                    Ok(())
321                };
322                set_stdio("stdin", 0, "r")?;
323                set_stdio("stdout", 1, "w")?;
324                set_stdio("stderr", 2, "w")?;
325
326                let io_open = io.get_attr("open", self)?;
327                self.builtins.set_attr("open", io_open, self)?;
328            }
329
330            Ok(importlib)
331        };
332
333        let res = essential_init();
334        let importlib = self.expect_pyresult(res, "essential initialization failed");
335
336        if self.state.settings.allow_external_library && cfg!(feature = "rustpython-compiler") {
337            if let Err(e) = import::init_importlib_package(self, importlib) {
338                eprintln!("importlib initialization failed. This is critical for many complicated packages.");
339                self.print_exception(e);
340            }
341        }
342
343        #[cfg(feature = "encodings")]
344        if cfg!(feature = "freeze-stdlib") || !self.state.settings.path_list.is_empty() {
345            if let Err(e) = self.import_encodings() {
346                eprintln!(
347                    "encodings initialization failed. Only utf-8 encoding will be supported."
348                );
349                self.print_exception(e);
350            }
351        } else {
352            // Here may not be the best place to give general `path_list` advice,
353            // but bare rustpython_vm::VirtualMachine users skipped proper settings must hit here while properly setup vm never enters here.
354            eprintln!(
355                "feature `encodings` is enabled but `settings.path_list` is empty. \
356                Please add the library path to `settings.path_list`. If you intended to disable the entire standard library (including the `encodings` feature), please also make sure to disable the `encodings` feature.\n\
357                Tip: You may also want to add `\"\"` to `settings.path_list` in order to enable importing from the current working directory."
358            );
359        }
360
361        self.initialized = true;
362    }
363
364    fn state_mut(&mut self) -> &mut PyGlobalState {
365        PyRc::get_mut(&mut self.state)
366            .expect("there should not be multiple threads while a user has a mut ref to a vm")
367    }
368
369    /// Can only be used in the initialization closure passed to [`Interpreter::with_init`]
370    pub fn add_native_module<S>(&mut self, name: S, module: stdlib::StdlibInitFunc)
371    where
372        S: Into<Cow<'static, str>>,
373    {
374        self.state_mut().module_inits.insert(name.into(), module);
375    }
376
377    pub fn add_native_modules<I>(&mut self, iter: I)
378    where
379        I: IntoIterator<Item = (Cow<'static, str>, stdlib::StdlibInitFunc)>,
380    {
381        self.state_mut().module_inits.extend(iter);
382    }
383
384    /// Can only be used in the initialization closure passed to [`Interpreter::with_init`]
385    pub fn add_frozen<I>(&mut self, frozen: I)
386    where
387        I: IntoIterator<Item = (&'static str, FrozenModule)>,
388    {
389        self.state_mut().frozen.extend(frozen);
390    }
391
392    /// Set the custom signal channel for the interpreter
393    pub fn set_user_signal_channel(&mut self, signal_rx: signal::UserSignalReceiver) {
394        self.signal_rx = Some(signal_rx);
395    }
396
397    pub fn run_code_obj(&self, code: PyRef<PyCode>, scope: Scope) -> PyResult {
398        let frame = Frame::new(code, scope, self.builtins.dict(), &[], self).into_ref(&self.ctx);
399        self.run_frame(frame)
400    }
401
402    #[cold]
403    pub fn run_unraisable(&self, e: PyBaseExceptionRef, msg: Option<String>, object: PyObjectRef) {
404        let sys_module = self.import("sys", 0).unwrap();
405        let unraisablehook = sys_module.get_attr("unraisablehook", self).unwrap();
406
407        let exc_type = e.class().to_owned();
408        let exc_traceback = e.traceback().to_pyobject(self); // TODO: actual traceback
409        let exc_value = e.into();
410        let args = stdlib::sys::UnraisableHookArgs {
411            exc_type,
412            exc_value,
413            exc_traceback,
414            err_msg: self.new_pyobj(msg),
415            object,
416        };
417        if let Err(e) = unraisablehook.call((args,), self) {
418            println!("{}", e.as_object().repr(self).unwrap().as_str());
419        }
420    }
421
422    #[inline(always)]
423    pub fn run_frame(&self, frame: FrameRef) -> PyResult {
424        match self.with_frame(frame, |f| f.run(self))? {
425            ExecutionResult::Return(value) => Ok(value),
426            _ => panic!("Got unexpected result from function"),
427        }
428    }
429
430    pub fn current_recursion_depth(&self) -> usize {
431        self.recursion_depth.get()
432    }
433
434    /// Used to run the body of a (possibly) recursive function. It will raise a
435    /// RecursionError if recursive functions are nested far too many times,
436    /// preventing a stack overflow.
437    pub fn with_recursion<R, F: FnOnce() -> PyResult<R>>(&self, _where: &str, f: F) -> PyResult<R> {
438        self.check_recursive_call(_where)?;
439        self.recursion_depth.set(self.recursion_depth.get() + 1);
440        let result = f();
441        self.recursion_depth.set(self.recursion_depth.get() - 1);
442        result
443    }
444
445    pub fn with_frame<R, F: FnOnce(FrameRef) -> PyResult<R>>(
446        &self,
447        frame: FrameRef,
448        f: F,
449    ) -> PyResult<R> {
450        self.with_recursion("", || {
451            self.frames.borrow_mut().push(frame.clone());
452            let result = f(frame);
453            // defer dec frame
454            let _popped = self.frames.borrow_mut().pop();
455            result
456        })
457    }
458
459    /// Returns a basic CompileOpts instance with options accurate to the vm. Used
460    /// as the CompileOpts for `vm.compile()`.
461    #[cfg(feature = "rustpython-codegen")]
462    pub fn compile_opts(&self) -> crate::compiler::CompileOpts {
463        crate::compiler::CompileOpts {
464            optimize: self.state.settings.optimize,
465        }
466    }
467
468    // To be called right before raising the recursion depth.
469    fn check_recursive_call(&self, _where: &str) -> PyResult<()> {
470        if self.recursion_depth.get() >= self.recursion_limit.get() {
471            Err(self.new_recursion_error(format!("maximum recursion depth exceeded {_where}")))
472        } else {
473            Ok(())
474        }
475    }
476
477    pub fn current_frame(&self) -> Option<Ref<FrameRef>> {
478        let frames = self.frames.borrow();
479        if frames.is_empty() {
480            None
481        } else {
482            Some(Ref::map(self.frames.borrow(), |frames| {
483                frames.last().unwrap()
484            }))
485        }
486    }
487
488    pub fn current_locals(&self) -> PyResult<ArgMapping> {
489        self.current_frame()
490            .expect("called current_locals but no frames on the stack")
491            .locals(self)
492    }
493
494    pub fn current_globals(&self) -> Ref<PyDictRef> {
495        let frame = self
496            .current_frame()
497            .expect("called current_globals but no frames on the stack");
498        Ref::map(frame, |f| &f.globals)
499    }
500
501    pub fn try_class(&self, module: &'static str, class: &'static str) -> PyResult<PyTypeRef> {
502        let class = self
503            .import(module, 0)?
504            .get_attr(class, self)?
505            .downcast()
506            .expect("not a class");
507        Ok(class)
508    }
509
510    pub fn class(&self, module: &'static str, class: &'static str) -> PyTypeRef {
511        let module = self
512            .import(module, 0)
513            .unwrap_or_else(|_| panic!("unable to import {module}"));
514
515        let class = module
516            .get_attr(class, self)
517            .unwrap_or_else(|_| panic!("module {module:?} has no class {class}"));
518        class.downcast().expect("not a class")
519    }
520
521    /// Call Python __import__ function without from_list.
522    /// Roughly equivalent to `import module_name` or `import top.submodule`.
523    ///
524    /// See also [`import_from`] for more advanced import.
525    /// See also [`rustpython_vm::import::import_source`] and other primitive import functions.
526    #[inline]
527    pub fn import<'a>(&self, module_name: impl AsPyStr<'a>, level: usize) -> PyResult {
528        let module_name = module_name.as_pystr(&self.ctx);
529        let from_list = PyTupleTyped::empty(self);
530        self.import_inner(module_name, from_list, level)
531    }
532
533    /// Call Python __import__ function caller with from_list.
534    /// Roughly equivalent to `from module_name import item1, item2` or `from top.submodule import item1, item2`
535    #[inline]
536    pub fn import_from<'a>(
537        &self,
538        module_name: impl AsPyStr<'a>,
539        from_list: PyTupleTyped<PyStrRef>,
540        level: usize,
541    ) -> PyResult {
542        let module_name = module_name.as_pystr(&self.ctx);
543        self.import_inner(module_name, from_list, level)
544    }
545
546    fn import_inner(
547        &self,
548        module: &Py<PyStr>,
549        from_list: PyTupleTyped<PyStrRef>,
550        level: usize,
551    ) -> PyResult {
552        // if the import inputs seem weird, e.g a package import or something, rather than just
553        // a straight `import ident`
554        let weird = module.as_str().contains('.') || level != 0 || !from_list.is_empty();
555
556        let cached_module = if weird {
557            None
558        } else {
559            let sys_modules = self.sys_module.get_attr("modules", self)?;
560            sys_modules.get_item(module, self).ok()
561        };
562
563        match cached_module {
564            Some(cached_module) => {
565                if self.is_none(&cached_module) {
566                    Err(self.new_import_error(
567                        format!("import of {module} halted; None in sys.modules"),
568                        module.to_owned(),
569                    ))
570                } else {
571                    Ok(cached_module)
572                }
573            }
574            None => {
575                let import_func = self
576                    .builtins
577                    .get_attr(identifier!(self, __import__), self)
578                    .map_err(|_| {
579                        self.new_import_error("__import__ not found".to_owned(), module.to_owned())
580                    })?;
581
582                let (locals, globals) = if let Some(frame) = self.current_frame() {
583                    (Some(frame.locals.clone()), Some(frame.globals.clone()))
584                } else {
585                    (None, None)
586                };
587                let from_list = from_list.to_pyobject(self);
588                import_func
589                    .call((module.to_owned(), globals, locals, from_list, level), self)
590                    .map_err(|exc| import::remove_importlib_frames(self, &exc))
591            }
592        }
593    }
594
595    pub fn extract_elements_with<T, F>(&self, value: &PyObject, func: F) -> PyResult<Vec<T>>
596    where
597        F: Fn(PyObjectRef) -> PyResult<T>,
598    {
599        // Extract elements from item, if possible:
600        let cls = value.class();
601        let list_borrow;
602        let slice = if cls.is(self.ctx.types.tuple_type) {
603            value.payload::<PyTuple>().unwrap().as_slice()
604        } else if cls.is(self.ctx.types.list_type) {
605            list_borrow = value.payload::<PyList>().unwrap().borrow_vec();
606            &list_borrow
607        } else {
608            return self.map_pyiter(value, func);
609        };
610        slice.iter().map(|obj| func(obj.clone())).collect()
611    }
612
613    pub fn map_iterable_object<F, R>(&self, obj: &PyObject, mut f: F) -> PyResult<PyResult<Vec<R>>>
614    where
615        F: FnMut(PyObjectRef) -> PyResult<R>,
616    {
617        match_class!(match obj {
618            ref l @ PyList => {
619                let mut i: usize = 0;
620                let mut results = Vec::with_capacity(l.borrow_vec().len());
621                loop {
622                    let elem = {
623                        let elements = &*l.borrow_vec();
624                        if i >= elements.len() {
625                            results.shrink_to_fit();
626                            return Ok(Ok(results));
627                        } else {
628                            elements[i].clone()
629                        }
630                        // free the lock
631                    };
632                    match f(elem) {
633                        Ok(result) => results.push(result),
634                        Err(err) => return Ok(Err(err)),
635                    }
636                    i += 1;
637                }
638            }
639            ref t @ PyTuple => Ok(t.iter().cloned().map(f).collect()),
640            // TODO: put internal iterable type
641            obj => {
642                Ok(self.map_pyiter(obj, f))
643            }
644        })
645    }
646
647    fn map_pyiter<F, R>(&self, value: &PyObject, mut f: F) -> PyResult<Vec<R>>
648    where
649        F: FnMut(PyObjectRef) -> PyResult<R>,
650    {
651        let iter = value.to_owned().get_iter(self)?;
652        let cap = match self.length_hint_opt(value.to_owned()) {
653            Err(e) if e.class().is(self.ctx.exceptions.runtime_error) => return Err(e),
654            Ok(Some(value)) => Some(value),
655            // Use a power of 2 as a default capacity.
656            _ => None,
657        };
658        // TODO: fix extend to do this check (?), see test_extend in Lib/test/list_tests.py,
659        // https://github.com/python/cpython/blob/v3.9.0/Objects/listobject.c#L922-L928
660        if let Some(cap) = cap {
661            if cap >= isize::MAX as usize {
662                return Ok(Vec::new());
663            }
664        }
665
666        let mut results = PyIterIter::new(self, iter.as_ref(), cap)
667            .map(|element| f(element?))
668            .collect::<PyResult<Vec<_>>>()?;
669        results.shrink_to_fit();
670        Ok(results)
671    }
672
673    pub fn get_attribute_opt<'a>(
674        &self,
675        obj: PyObjectRef,
676        attr_name: impl AsPyStr<'a>,
677    ) -> PyResult<Option<PyObjectRef>> {
678        let attr_name = attr_name.as_pystr(&self.ctx);
679        match obj.get_attr_inner(attr_name, self) {
680            Ok(attr) => Ok(Some(attr)),
681            Err(e) if e.fast_isinstance(self.ctx.exceptions.attribute_error) => Ok(None),
682            Err(e) => Err(e),
683        }
684    }
685
686    pub fn set_attribute_error_context(
687        &self,
688        exc: &PyBaseExceptionRef,
689        obj: PyObjectRef,
690        name: PyStrRef,
691    ) {
692        if exc.class().is(self.ctx.exceptions.attribute_error) {
693            let exc = exc.as_object();
694            exc.set_attr("name", name, self).unwrap();
695            exc.set_attr("obj", obj, self).unwrap();
696        }
697    }
698
699    // get_method should be used for internal access to magic methods (by-passing
700    // the full getattribute look-up.
701    pub fn get_method_or_type_error<F>(
702        &self,
703        obj: PyObjectRef,
704        method_name: &'static PyStrInterned,
705        err_msg: F,
706    ) -> PyResult
707    where
708        F: FnOnce() -> String,
709    {
710        let method = obj
711            .class()
712            .get_attr(method_name)
713            .ok_or_else(|| self.new_type_error(err_msg()))?;
714        self.call_if_get_descriptor(&method, obj)
715    }
716
717    // TODO: remove + transfer over to get_special_method
718    pub(crate) fn get_method(
719        &self,
720        obj: PyObjectRef,
721        method_name: &'static PyStrInterned,
722    ) -> Option<PyResult> {
723        let method = obj.get_class_attr(method_name)?;
724        Some(self.call_if_get_descriptor(&method, obj))
725    }
726
727    pub(crate) fn get_str_method(&self, obj: PyObjectRef, method_name: &str) -> Option<PyResult> {
728        let method_name = self.ctx.interned_str(method_name)?;
729        self.get_method(obj, method_name)
730    }
731
732    #[inline]
733    /// Checks for triggered signals and calls the appropriate handlers. A no-op on
734    /// platforms where signals are not supported.
735    pub fn check_signals(&self) -> PyResult<()> {
736        #[cfg(not(target_arch = "wasm32"))]
737        {
738            crate::signal::check_signals(self)
739        }
740        #[cfg(target_arch = "wasm32")]
741        {
742            Ok(())
743        }
744    }
745
746    pub(crate) fn push_exception(&self, exc: Option<PyBaseExceptionRef>) {
747        let mut excs = self.exceptions.borrow_mut();
748        let prev = std::mem::take(&mut *excs);
749        excs.prev = Some(Box::new(prev));
750        excs.exc = exc
751    }
752
753    pub(crate) fn pop_exception(&self) -> Option<PyBaseExceptionRef> {
754        let mut excs = self.exceptions.borrow_mut();
755        let cur = std::mem::take(&mut *excs);
756        *excs = *cur.prev.expect("pop_exception() without nested exc stack");
757        cur.exc
758    }
759
760    pub(crate) fn take_exception(&self) -> Option<PyBaseExceptionRef> {
761        self.exceptions.borrow_mut().exc.take()
762    }
763
764    pub(crate) fn current_exception(&self) -> Option<PyBaseExceptionRef> {
765        self.exceptions.borrow().exc.clone()
766    }
767
768    pub(crate) fn set_exception(&self, exc: Option<PyBaseExceptionRef>) {
769        // don't be holding the RefCell guard while __del__ is called
770        let prev = std::mem::replace(&mut self.exceptions.borrow_mut().exc, exc);
771        drop(prev);
772    }
773
774    pub(crate) fn contextualize_exception(&self, exception: &PyBaseExceptionRef) {
775        if let Some(context_exc) = self.topmost_exception() {
776            if !context_exc.is(exception) {
777                let mut o = context_exc.clone();
778                while let Some(context) = o.context() {
779                    if context.is(exception) {
780                        o.set_context(None);
781                        break;
782                    }
783                    o = context;
784                }
785                exception.set_context(Some(context_exc))
786            }
787        }
788    }
789
790    pub(crate) fn topmost_exception(&self) -> Option<PyBaseExceptionRef> {
791        let excs = self.exceptions.borrow();
792        let mut cur = &*excs;
793        loop {
794            if let Some(exc) = &cur.exc {
795                return Some(exc.clone());
796            }
797            cur = cur.prev.as_deref()?;
798        }
799    }
800
801    pub fn handle_exit_exception(&self, exc: PyBaseExceptionRef) -> u8 {
802        if exc.fast_isinstance(self.ctx.exceptions.system_exit) {
803            let args = exc.args();
804            let msg = match args.as_slice() {
805                [] => return 0,
806                [arg] => match_class!(match arg {
807                    ref i @ PyInt => {
808                        use num_traits::cast::ToPrimitive;
809                        return i.as_bigint().to_u8().unwrap_or(0);
810                    }
811                    arg => {
812                        if self.is_none(arg) {
813                            return 0;
814                        } else {
815                            arg.str(self).ok()
816                        }
817                    }
818                }),
819                _ => args.as_object().repr(self).ok(),
820            };
821            if let Some(msg) = msg {
822                let stderr = stdlib::sys::PyStderr(self);
823                writeln!(stderr, "{msg}");
824            }
825            1
826        } else if exc.fast_isinstance(self.ctx.exceptions.keyboard_interrupt) {
827            #[allow(clippy::if_same_then_else)]
828            {
829                self.print_exception(exc);
830                #[cfg(unix)]
831                {
832                    let action = SigAction::new(
833                        nix::sys::signal::SigHandler::SigDfl,
834                        SaFlags::SA_ONSTACK,
835                        SigSet::empty(),
836                    );
837                    let result = unsafe { sigaction(SIGINT, &action) };
838                    if result.is_ok() {
839                        self.flush_std();
840                        kill(getpid(), SIGINT).expect("Expect to be killed.");
841                    }
842
843                    (libc::SIGINT as u8) + 128u8
844                }
845                #[cfg(not(unix))]
846                {
847                    1
848                }
849            }
850        } else {
851            self.print_exception(exc);
852            1
853        }
854    }
855
856    #[doc(hidden)]
857    pub fn __module_set_attr(
858        &self,
859        module: &Py<PyModule>,
860        attr_name: &'static PyStrInterned,
861        attr_value: impl Into<PyObjectRef>,
862    ) -> PyResult<()> {
863        let val = attr_value.into();
864        module
865            .as_object()
866            .generic_setattr(attr_name, PySetterValue::Assign(val), self)
867    }
868
869    pub fn insert_sys_path(&self, obj: PyObjectRef) -> PyResult<()> {
870        let sys_path = self.sys_module.get_attr("path", self).unwrap();
871        self.call_method(&sys_path, "insert", (0, obj))?;
872        Ok(())
873    }
874
875    pub fn run_module(&self, module: &str) -> PyResult<()> {
876        let runpy = self.import("runpy", 0)?;
877        let run_module_as_main = runpy.get_attr("_run_module_as_main", self)?;
878        run_module_as_main.call((module,), self)?;
879        Ok(())
880    }
881}
882
883impl AsRef<Context> for VirtualMachine {
884    fn as_ref(&self) -> &Context {
885        &self.ctx
886    }
887}
888
889fn core_frozen_inits() -> impl Iterator<Item = (&'static str, FrozenModule)> {
890    let iter = std::iter::empty();
891    macro_rules! ext_modules {
892        ($iter:ident, $($t:tt)*) => {
893            let $iter = $iter.chain(py_freeze!($($t)*));
894        };
895    }
896
897    // keep as example but use file one now
898    // ext_modules!(
899    //     iter,
900    //     source = "initialized = True; print(\"Hello world!\")\n",
901    //     module_name = "__hello__",
902    // );
903
904    // Python modules that the vm calls into, but are not actually part of the stdlib. They could
905    // in theory be implemented in Rust, but are easiest to do in Python for one reason or another.
906    // Includes _importlib_bootstrap and _importlib_bootstrap_external
907    ext_modules!(
908        iter,
909        dir = "./Lib/python_builtins",
910        crate_name = "rustpython_compiler_core"
911    );
912
913    // core stdlib Python modules that the vm calls into, but are still used in Python
914    // application code, e.g. copyreg
915    // FIXME: Initializing core_modules here results duplicated frozen module generation for core_modules.
916    // We need a way to initialize this modules for both `Interpreter::without_stdlib()` and `InterpreterConfig::new().init_stdlib().interpreter()`
917    // #[cfg(not(feature = "freeze-stdlib"))]
918    ext_modules!(
919        iter,
920        dir = "./Lib/core_modules",
921        crate_name = "rustpython_compiler_core"
922    );
923
924    iter
925}
926
927#[test]
928fn test_nested_frozen() {
929    use rustpython_vm as vm;
930
931    vm::Interpreter::with_init(Default::default(), |vm| {
932        // vm.add_native_modules(rustpython_stdlib::get_module_inits());
933        vm.add_frozen(rustpython_vm::py_freeze!(dir = "../extra_tests/snippets"));
934    })
935    .enter(|vm| {
936        let scope = vm.new_scope_with_builtins();
937
938        let source = "from dir_module.dir_module_inner import value2";
939        let code_obj = vm
940            .compile(source, vm::compiler::Mode::Exec, "<embedded>".to_owned())
941            .map_err(|err| vm.new_syntax_error(&err, Some(source)))
942            .unwrap();
943
944        if let Err(e) = vm.run_code_obj(code_obj, scope) {
945            vm.print_exception(e);
946            panic!();
947        }
948    })
949}