Skip to main content

rustpython_vm/vm/
crossinterp.rs

1//! Cross-interpreter object sharing and exception snapshots (XI data).
2//!
3//! Interpreters do not share `PyObject` graphs. Shareable values are converted
4//! to an interpreter-neutral payload and rebuilt in the destination.
5// cspell:ignore unpickler
6
7#[cfg(not(feature = "threading"))]
8use crate::protocol::BufferDescriptor;
9use crate::{
10    AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromBorrowedObject,
11    VirtualMachine,
12    builtins::{
13        PyBaseException, PyBaseExceptionRef, PyBytes, PyCode, PyDict, PyFloat, PyFunction, PyInt,
14        PyMemoryView, PyStr, PyStrInterned, PyTuple,
15    },
16    bytecode::{
17        BorrowedConstant, CodeFlags, Constant, Instruction,
18        oparg::{OpArg, OpArgState},
19    },
20    protocol::PyBuffer,
21};
22use core::sync::atomic::{AtomicBool, Ordering};
23use num_traits::ToPrimitive;
24use rustpython_common::wtf8::Wtf8Buf;
25
26static MEMORYVIEW_REGISTERED: AtomicBool = AtomicBool::new(false);
27
28/// `register_memoryview_xid`: the builtin memoryview gets its getdata function
29/// only once `_interpreters` has been imported.
30pub(crate) fn register_memoryview_xid() {
31    MEMORYVIEW_REGISTERED.store(true, Ordering::Relaxed);
32}
33
34fn memoryview_registered() -> bool {
35    MEMORYVIEW_REGISTERED.load(Ordering::Relaxed)
36}
37
38/// Unbound-item ops (`UNBOUND_REMOVE` / `UNBOUND_ERROR` / `UNBOUND_REPLACE`).
39#[repr(i32)]
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum UnboundOp {
42    Remove = 1,
43    Error = 2,
44    Replace = 3,
45}
46
47impl UnboundOp {
48    #[must_use]
49    pub const fn from_i32(v: i32) -> Option<Self> {
50        Some(match v {
51            1 => Self::Remove,
52            2 => Self::Error,
53            3 => Self::Replace,
54            _ => return None,
55        })
56    }
57
58    #[must_use]
59    pub const fn as_i32(self) -> i32 {
60        self as i32
61    }
62}
63
64/// `xidata_fallback_t`: what to do with an object that has no XI data support.
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum Fallback {
67    /// `_PyXIDATA_XIDATA_ONLY`: raise `NotShareableError`.
68    XidataOnly = 0,
69    /// `_PyXIDATA_FULL_FALLBACK`: fall back to stateless functions, then pickle.
70    Full = 1,
71}
72
73impl Fallback {
74    #[must_use]
75    pub const fn from_i32(v: i32) -> Option<Self> {
76        Some(match v {
77            0 => Self::XidataOnly,
78            1 => Self::Full,
79            _ => return None,
80        })
81    }
82
83    #[must_use]
84    pub const fn as_i32(self) -> i32 {
85        self as i32
86    }
87}
88
89/// A queue reference held by cross-interpreter data. Cloning takes another
90/// reference and dropping releases one.
91pub struct SharedQueueId {
92    qid: i64,
93}
94
95impl SharedQueueId {
96    fn new(qid: i64) -> Option<Self> {
97        if crate::stdlib::_interpqueues::queue_xid_incref(qid) {
98            Some(Self { qid })
99        } else {
100            None
101        }
102    }
103}
104
105impl Clone for SharedQueueId {
106    fn clone(&self) -> Self {
107        let held = crate::stdlib::_interpqueues::queue_xid_incref(self.qid);
108        debug_assert!(held);
109        Self { qid: self.qid }
110    }
111}
112
113impl Drop for SharedQueueId {
114    fn drop(&mut self) {
115        crate::stdlib::_interpqueues::queue_xid_decref(self.qid);
116    }
117}
118
119/// Interpreter-neutral payload that can be materialized in any interpreter
120/// sharing the process-wide [`crate::Context`].
121#[derive(Clone)]
122pub enum SharedValue {
123    None,
124    Bool(bool),
125    Int(isize),
126    Float(f64),
127    Bytes(Vec<u8>),
128    Str(Wtf8Buf),
129    Tuple(Vec<Self>),
130    Channel {
131        cid: i64,
132        end: i32,
133    },
134    Queue(SharedQueueId),
135    #[cfg(feature = "threading")]
136    Buffer(PyBuffer),
137    #[cfg(not(feature = "threading"))]
138    Buffer(Vec<u8>, BufferDescriptor),
139    /// Marshalled code object.
140    Code(Vec<u8>),
141    /// Marshalled code of a stateless function; rebuilt against `__main__`.
142    Function(Vec<u8>),
143    /// `pickle.dumps` output, used by [`Fallback::Full`].
144    Pickled(PickledData),
145}
146
147/// Pickled bytes and the originating script, without interpreter-owned references.
148#[derive(Debug, Clone)]
149pub struct PickledData {
150    bytes: Vec<u8>,
151    mainfile: Option<String>,
152}
153
154impl SharedValue {
155    pub fn from_object(obj: &PyObject, fallback: Fallback, vm: &VirtualMachine) -> PyResult<Self> {
156        match Self::basic_from_object(obj, fallback, vm) {
157            Ok(v) => Ok(v),
158            Err(exc) => {
159                if fallback == Fallback::XidataOnly {
160                    return Err(exc);
161                }
162                if obj.downcastable::<PyFunction>()
163                    && let Ok(v) = Self::from_function(obj, vm)
164                {
165                    return Ok(v);
166                }
167                // We could try marshal here but we don't for now.
168                match pickle_dumps(obj, vm) {
169                    Ok(data) => Ok(Self::Pickled(data)),
170                    // Raise the original exception.
171                    Err(_) => Err(exc),
172                }
173            }
174        }
175    }
176
177    /// `_get_xidata`: the registered "getdata" functions, i.e. no pickle
178    /// fallback. Both a missing registration and a failing conversion are
179    /// reported under `obj`, the latter keeping its own failure as the cause.
180    fn basic_from_object(
181        obj: &PyObject,
182        fallback: Fallback,
183        vm: &VirtualMachine,
184    ) -> PyResult<Self> {
185        match Self::getdata(obj, fallback, vm) {
186            Some(res) => res.map_err(|cause| not_shareable(vm, obj, Some(cause))),
187            None => Err(not_shareable(vm, obj, None)),
188        }
189    }
190
191    /// `lookup_getdata` followed by the call itself. `None` when the type has
192    /// no registered getdata function.
193    fn getdata(obj: &PyObject, fallback: Fallback, vm: &VirtualMachine) -> Option<PyResult<Self>> {
194        if vm.is_none(obj) {
195            return Some(Ok(Self::None));
196        }
197        let cls = obj.class();
198        if cls.is(vm.ctx.types.bool_type) {
199            return Some(obj.to_owned().is_true(vm).map(Self::Bool));
200        }
201        if cls.is(vm.ctx.types.int_type) {
202            let n = obj.downcast_ref::<PyInt>().unwrap();
203            // `_long_shared`: the size of shareable ints is bounded by
204            // sys.maxsize.
205            return Some(
206                n.as_bigint()
207                    .to_isize()
208                    .map(Self::Int)
209                    .ok_or_else(|| vm.new_overflow_error("try sending as bytes")),
210            );
211        }
212        if cls.is(vm.ctx.types.float_type) {
213            let f = obj.downcast_ref::<PyFloat>().unwrap();
214            return Some(Ok(Self::Float(f.to_f64())));
215        }
216        if cls.is(vm.ctx.types.bytes_type) {
217            let b = obj.downcast_ref::<PyBytes>().unwrap();
218            return Some(Ok(Self::Bytes(b.as_bytes().to_vec())));
219        }
220        if cls.is(vm.ctx.types.str_type) {
221            let s = obj.downcast_ref::<PyStr>().unwrap();
222            return Some(Ok(Self::Str(s.as_wtf8().to_owned())));
223        }
224        if cls.is(vm.ctx.types.tuple_type) {
225            let t = obj.downcast_ref::<PyTuple>().unwrap();
226            return Some(
227                t.as_slice()
228                    .iter()
229                    .map(|item| {
230                        vm.with_recursion("while sharing a tuple", || {
231                            Self::from_object(item, fallback, vm)
232                        })
233                    })
234                    .collect::<PyResult<Vec<_>>>()
235                    .map(Self::Tuple),
236            );
237        }
238        // `_pybuffer_shared`, registered for the builtin memoryview by
239        // `_interpreters` and never unregistered.
240        if cls.is(vm.ctx.types.memoryview_type) && memoryview_registered() {
241            return Some(Self::from_buffer_object(obj, vm));
242        }
243        // Registered by the `_interpchannels` module for ChannelID.
244        if let Some(ch) = crate::stdlib::_interpchannels::channel_id_parts(obj) {
245            return Some(Ok(Self::Channel {
246                cid: ch.0,
247                end: ch.1,
248            }));
249        }
250        if let Some(qid) = crate::stdlib::_interpqueues::queue_id_from_object(obj, vm) {
251            return match qid {
252                Ok(qid) => SharedQueueId::new(qid).map(|qid| Ok(Self::Queue(qid))),
253                Err(exc) => Some(Err(exc)),
254            };
255        }
256        None
257    }
258
259    /// `_PyFunction_GetXIData`: only stateless functions are shareable.
260    fn from_function(obj: &PyObject, vm: &VirtualMachine) -> PyResult<Self> {
261        let func = obj.downcast_ref::<PyFunction>().ok_or_else(|| {
262            not_shareable_error(
263                vm,
264                format!("expected a function, got {}", render_repr(obj, vm)),
265            )
266        })?;
267        verify_stateless_function(func, vm).map_err(|cause| {
268            not_shareable_error_from(vm, "only stateless functions are shareable", cause)
269        })?;
270        let code = (*func.code).to_owned();
271        Ok(Self::Function(marshal_dumps(code.as_object(), vm)?))
272    }
273
274    /// `_interp_call_pack`'s handling of the callable: a stateless function
275    /// travels as its code, anything else is pickled.
276    pub fn from_callable(obj: &PyObject, vm: &VirtualMachine) -> PyResult<Self> {
277        match Self::from_function(obj, vm) {
278            Ok(v) => Ok(v),
279            Err(exc) => match pickle_dumps(obj, vm) {
280                Ok(data) => Ok(Self::Pickled(data)),
281                // Raise the exception from the stateless check.
282                Err(_) => Err(exc),
283            },
284        }
285    }
286
287    /// `_PyCode_GetXIData`.
288    pub fn from_code(code: &Py<PyCode>, vm: &VirtualMachine) -> PyResult<Self> {
289        Ok(Self::Code(marshal_dumps(code.as_object(), vm)?))
290    }
291
292    pub fn from_buffer_object(obj: &PyObject, vm: &VirtualMachine) -> PyResult<Self> {
293        let view = PyMemoryView::from_object(obj, vm)?;
294        #[cfg(feature = "threading")]
295        {
296            Ok(Self::Buffer(view.clone_buffer()))
297        }
298        #[cfg(not(feature = "threading"))]
299        {
300            let buf = view.clone_buffer();
301            let bytes = buf.obj_bytes().to_vec();
302            Ok(Self::Buffer(bytes, buf.desc.clone()))
303        }
304    }
305
306    pub fn into_object(self, vm: &VirtualMachine) -> PyResult {
307        match self {
308            Self::None => Ok(vm.ctx.none()),
309            Self::Bool(b) => Ok(vm.ctx.new_bool(b).into()),
310            Self::Int(n) => Ok(vm.ctx.new_int(n).into()),
311            Self::Float(f) => Ok(vm.ctx.new_float(f).into()),
312            Self::Bytes(b) => Ok(vm.ctx.new_bytes(b).into()),
313            Self::Str(s) => Ok(vm.ctx.new_str(s).into()),
314            Self::Tuple(items) => {
315                let mut els = Vec::with_capacity(items.len());
316                for item in items {
317                    els.push(item.into_object(vm)?);
318                }
319                Ok(vm.ctx.new_tuple(els).into())
320            }
321            Self::Channel { cid, end } => {
322                crate::stdlib::_interpchannels::channel_id_from_parts(cid, end, false, false, vm)
323            }
324            Self::Queue(queue) => crate::stdlib::_interpqueues::queue_from_xid(queue.qid, vm),
325            #[cfg(feature = "threading")]
326            Self::Buffer(buffer) => memoryview_from_buffer(buffer, vm),
327            #[cfg(not(feature = "threading"))]
328            Self::Buffer(bytes, desc) => {
329                use crate::protocol::VecBuffer;
330                let buffer = VecBuffer::from(bytes)
331                    .into_ref(&vm.ctx)
332                    .into_pybuffer_with_descriptor(desc);
333                memoryview_from_buffer(buffer, vm)
334            }
335            Self::Code(data) => marshal_loads(&data, vm),
336            Self::Function(data) => {
337                let code = marshal_loads(&data, vm)?
338                    .downcast::<PyCode>()
339                    .map_err(|_| vm.new_type_error("expected code"))?;
340                // Stateless functions have no globals, so `__main__` is used,
341                // just like for builtins such as exec().
342                let globals = vm.main_namespace()?;
343                Ok(PyFunction::new(code, globals, vm)?.into_pyobject(vm))
344            }
345            Self::Pickled(data) => pickle_loads(&data, vm),
346        }
347    }
348}
349
350fn marshal_dumps(obj: &PyObject, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
351    let dumps = vm.import("marshal", 0)?.get_attr("dumps", vm)?;
352    let bytes = dumps.call((obj.to_owned(),), vm)?;
353    let bytes = bytes
354        .downcast::<PyBytes>()
355        .map_err(|_| vm.new_type_error("marshal.dumps() did not return bytes"))?;
356    Ok(bytes.as_bytes().to_vec())
357}
358
359fn marshal_loads(data: &[u8], vm: &VirtualMachine) -> PyResult {
360    let loads = vm.import("marshal", 0)?.get_attr("loads", vm)?;
361    loads.call((vm.ctx.new_bytes(data.to_vec()),), vm)
362}
363
364/// `_memoryview_from_xid`.
365fn memoryview_from_buffer(buffer: PyBuffer, vm: &VirtualMachine) -> PyResult {
366    let view = crate::stdlib::_interpreters::xibufferview_from_buffer(buffer, vm);
367    let mv = PyMemoryView::from_object(&view, vm)?;
368    Ok(mv.into_pyobject(vm))
369}
370
371fn pickle_dumps(obj: &PyObject, vm: &VirtualMachine) -> PyResult<PickledData> {
372    let dumps = vm.import("pickle", 0)?.get_attr("dumps", vm)?;
373    let bytes = dumps.call((obj.to_owned(),), vm)?;
374    let bytes = bytes
375        .downcast::<PyBytes>()
376        .map_err(|_| vm.new_type_error("pickle.dumps() did not return bytes"))?;
377    // Missing or non-string __file__ does not prevent ordinary pickling.
378    let mainfile = vm
379        .main_namespace()
380        .ok()
381        .and_then(|ns| ns.get_item_opt("__file__", vm).ok().flatten())
382        .and_then(|file| {
383            file.downcast_ref::<PyStr>()
384                .and_then(|s| s.to_str().map(str::to_owned))
385        });
386    Ok(PickledData {
387        bytes: bytes.as_bytes().to_vec(),
388        mainfile,
389    })
390}
391
392fn pickle_loads(data: &PickledData, vm: &VirtualMachine) -> PyResult {
393    let loads = vm.import("pickle", 0)?.get_attr("loads", vm)?;
394    let bytes = vm.ctx.new_bytes(data.bytes.clone());
395    let cause = match loads.call((bytes.clone(),), vm) {
396        Ok(obj) => return Ok(obj),
397        Err(cause) => cause,
398    };
399    let missing_main_attr = cause.fast_isinstance(vm.ctx.exceptions.attribute_error)
400        && cause
401            .args()
402            .as_slice()
403            .first()
404            .and_then(|arg| arg.downcast_ref::<PyStr>())
405            .is_some_and(|msg| {
406                msg.as_bytes()
407                    .starts_with(b"module '__main__' has no attribute '")
408            });
409    if missing_main_attr && let Some(mainfile) = &data.mainfile {
410        let retry = || -> PyResult {
411            let isolated = isolated_main(mainfile, vm)?;
412            let file = vm
413                .import("io", 0)?
414                .get_attr("BytesIO", vm)?
415                .call((bytes,), vm)?;
416            let unpickler = vm
417                .import("pickle", 0)?
418                .get_attr("Unpickler", vm)?
419                .call((file,), vm)?;
420            let original = unpickler.get_attr("find_class", vm)?;
421            let state = vm.ctx.new_tuple(vec![isolated, original]);
422            const FIND_CLASS: crate::function::PyMethodDef =
423                crate::function::PyMethodDef::new_const(
424                    "find_class",
425                    find_main_class,
426                    crate::function::PyMethodFlags::empty(),
427                    crate::function::ItemDoc::NONE,
428                );
429            unpickler.set_attr(
430                "find_class",
431                FIND_CLASS.build_bound_function(&vm.ctx, state.into()),
432                vm,
433            )?;
434            // Resolve globals on this unpickler only: replacing sys.modules
435            // would let concurrent consumers observe or restore the wrong module.
436            let _clear = scopeguard::guard((), |()| {
437                let _ = unpickler.del_attr("find_class", vm);
438            });
439            vm.call_method(&unpickler, "load", ())
440        };
441        if let Ok(obj) = retry() {
442            return Ok(obj);
443        }
444    }
445    // Like CPython, retain the original unpickling error if the retry fails.
446    Err(not_shareable_error_from(
447        vm,
448        "object could not be unpickled",
449        cause,
450    ))
451}
452
453fn find_main_class(
454    state: PyRef<PyTuple>,
455    module: PyRef<PyStr>,
456    name: PyRef<PyStr>,
457    vm: &VirtualMachine,
458) -> PyResult {
459    let [isolated, original] = state.as_slice() else {
460        unreachable!()
461    };
462    if module.as_bytes() != b"__main__" {
463        return original.call((module, name), vm);
464    }
465    vm.audit("pickle.find_class", || (module, name.clone()))?;
466    let mut value = isolated.clone();
467    if let Some(name) = name.to_str() {
468        for part in name.split('.') {
469            value = value.get_attr(&vm.ctx.new_str(part), vm)?;
470        }
471        Ok(value)
472    } else {
473        value.get_attr(&name, vm)
474    }
475}
476
477fn isolated_main(mainfile: &str, vm: &VirtualMachine) -> PyResult {
478    let module = vm.import("_interpreters", 0)?;
479    let state = module
480        .dict()
481        .ok_or_else(|| vm.new_runtime_error("missing _interpreters namespace"))?;
482    if let Some(loaded) = state.get_item_opt("_cached_main", vm)? {
483        return Ok(loaded);
484    }
485    // Use the import machinery's reentrant, deadlock-aware per-module lock
486    // while executing the script, including detaching threads that wait for it.
487    let manager = vm
488        .import("_frozen_importlib", 0)?
489        .get_attr("_ModuleLockManager", vm)?
490        .call(("<RustPython isolated __main__>",), vm)?;
491    vm.call_method(&manager, "__enter__", ())?;
492    let _unlock = scopeguard::guard((), |()| {
493        let _ = vm.call_method(
494            &manager,
495            "__exit__",
496            (vm.ctx.none(), vm.ctx.none(), vm.ctx.none()),
497        );
498    });
499    if let Some(loaded) = state.get_item_opt("_cached_main", vm)? {
500        return Ok(loaded);
501    }
502    let loaded = vm.new_module("__main__", vm.ctx.new_dict(), None);
503    let run_path = vm.import("runpy", 0)?.get_attr("run_path", vm)?;
504    // Avoid executing the script's `if __name__ == '__main__'` entry point.
505    let namespace = run_path.call((mainfile, vm.ctx.none(), "<fake __main__>"), vm)?;
506    loaded.dict().merge_object(namespace, vm)?;
507    // The native module owns this cache so normal module teardown releases it.
508    state.setdefault(vm.ctx.new_str("_cached_main").into(), loaded.into(), vm)
509}
510
511fn not_shareable_error(vm: &VirtualMachine, msg: impl Into<String>) -> PyBaseExceptionRef {
512    crate::stdlib::_interpreters::not_shareable_error(vm, msg)
513}
514
515/// `set_notshareableerror`: the failure that led here becomes the cause.
516fn not_shareable_error_from(
517    vm: &VirtualMachine,
518    msg: impl Into<String>,
519    cause: PyBaseExceptionRef,
520) -> PyBaseExceptionRef {
521    let exc = not_shareable_error(vm, msg);
522    exc.set_cause(Some(cause));
523    exc
524}
525
526/// `_set_xid_lookup_failure` without a message of its own.
527fn not_shareable(
528    vm: &VirtualMachine,
529    obj: &PyObject,
530    cause: Option<PyBaseExceptionRef>,
531) -> PyBaseExceptionRef {
532    let msg = format!(
533        "{} does not support cross-interpreter data",
534        render_repr(obj, vm)
535    );
536    match cause {
537        Some(cause) => not_shareable_error_from(vm, msg, cause),
538        None => not_shareable_error(vm, msg),
539    }
540}
541
542fn render_repr(obj: &PyObject, vm: &VirtualMachine) -> String {
543    obj.repr(vm)
544        .map_or_else(|_| obj.class().name().to_string(), |s| s.to_string())
545}
546
547/// Exact-type shareable check used by `_interpreters.is_shareable`.
548///
549/// This mirrors `_PyObject_CheckXIData`: only the type's registration is
550/// consulted, so a tuple is "shareable" even when its items are not.
551pub fn is_shareable(obj: &PyObject, vm: &VirtualMachine) -> bool {
552    if vm.is_none(obj) {
553        return true;
554    }
555    let cls = obj.class();
556    cls.is(vm.ctx.types.bool_type)
557        || cls.is(vm.ctx.types.int_type)
558        || cls.is(vm.ctx.types.float_type)
559        || cls.is(vm.ctx.types.bytes_type)
560        || cls.is(vm.ctx.types.str_type)
561        || cls.is(vm.ctx.types.tuple_type)
562        || (cls.is(vm.ctx.types.memoryview_type) && memoryview_registered())
563        || crate::stdlib::_interpchannels::channel_id_parts(obj).is_some()
564        || crate::stdlib::_interpqueues::is_external_queue(obj, vm)
565}
566
567/// Encode a namespace key as UTF-8, rejecting non-strings and lone surrogates.
568pub fn utf8_key<'a>(key: &'a PyObject, vm: &'a VirtualMachine) -> PyResult<&'a str> {
569    let s = key
570        .downcast_ref::<PyStr>()
571        .ok_or_else(|| vm.new_type_error("bad argument type for built-in operation"))?;
572    s.as_wtf8().as_str().map_err(|_| {
573        vm.new_unicode_encode_error(
574            vm.ctx.new_str("utf-8"),
575            s.to_owned(),
576            0,
577            s.char_len(),
578            vm.ctx.new_str("surrogates not allowed"),
579        )
580    })
581}
582
583/// `_convert_exc_to_TracebackException` followed by `_format_TracebackException`.
584fn format_traceback_exception(exc: &Py<PyBaseException>, vm: &VirtualMachine) -> PyResult<String> {
585    let create = vm
586        .import("traceback", 0)?
587        .get_attr("TracebackException", vm)?
588        .get_attr("from_exception", vm)?;
589    let kwargs: crate::function::KwArgs = [
590        (
591            "save_exc_type".to_owned(),
592            vm.ctx.false_value.clone().into(),
593        ),
594        ("lookup_lines".to_owned(), vm.ctx.false_value.clone().into()),
595    ]
596    .into_iter()
597    .collect();
598    let tbexc = create.call(
599        crate::function::FuncArgs::new(vec![exc.as_object().to_owned()], kwargs),
600        vm,
601    )?;
602    let lines = vm.call_method(&tbexc, "format", ())?;
603    let empty: PyObjectRef = vm.ctx.empty_str.to_owned().into();
604    let mut formatted = vm
605        .call_method(&empty, "join", (lines,))?
606        .str(vm)?
607        .to_string();
608    // Remove a trailing newline if needed.
609    if formatted.ends_with('\n') {
610        formatted.pop();
611    }
612    Ok(formatted)
613}
614
615/// `_PyXI_excinfo_format`, with `_excinfo_normalize_type` folded in: the
616/// `builtins` and `__main__` modules are left out of the name.
617fn format_exc_snapshot(module: &str, qualname: &str, msg: Option<&str>) -> String {
618    let name = if module == "builtins" || module == "__main__" {
619        qualname.to_owned()
620    } else {
621        format!("{module}.{qualname}")
622    };
623    match msg {
624        Some(msg) => format!("{name}: {msg}"),
625        None => name,
626    }
627}
628
629/// Snapshot of an exception that can be rebuilt in another interpreter.
630pub struct ExcInfo {
631    pub type_name: String,
632    pub type_qualname: String,
633    pub type_module: String,
634    pub msg: Option<String>,
635    pub formatted: String,
636    pub errdisplay: String,
637}
638
639impl ExcInfo {
640    pub fn capture(exc: &Py<PyBaseException>, vm: &VirtualMachine) -> Self {
641        let cls = exc.class();
642        let type_name = cls.name().to_owned();
643        let type_qualname = cls
644            .as_object()
645            .get_attr("__qualname__", vm)
646            .ok()
647            .and_then(|o| o.str(vm).ok())
648            .map_or_else(|| type_name.clone(), |s| s.to_string());
649        let type_module = cls
650            .as_object()
651            .get_attr("__module__", vm)
652            .ok()
653            .and_then(|o| o.str(vm).ok())
654            .map_or_else(|| "builtins".to_owned(), |s| s.to_string());
655        let msg = exc.as_object().str(vm).ok().map(|s| s.to_string());
656        let formatted = format_exc_snapshot(&type_module, &type_qualname, msg.as_deref());
657        let errdisplay = format_traceback_exception(exc, vm).unwrap_or_else(|_| {
658            let mut buf = String::new();
659            let _ = vm.write_exception(&mut buf, exc);
660            buf
661        });
662        Self {
663            type_name,
664            type_qualname,
665            type_module,
666            msg,
667            formatted,
668            errdisplay,
669        }
670    }
671
672    /// `_PyXI_ApplyError` for `_PyXI_ERR_NOT_SHAREABLE`: the snapshot is
673    /// re-raised as a plain `Exception` and becomes the cause.
674    pub fn into_not_shareable(
675        self,
676        msg: Option<String>,
677        vm: &VirtualMachine,
678    ) -> PyBaseExceptionRef {
679        let cause = vm.new_exception_msg(
680            vm.ctx.exceptions.exception_type.to_owned(),
681            self.errdisplay.into(),
682        );
683        let msg =
684            msg.unwrap_or_else(|| "object does not support cross-interpreter data".to_owned());
685        not_shareable_error_from(vm, msg, cause)
686    }
687
688    pub fn into_namespace(self, vm: &VirtualMachine) -> PyObjectRef {
689        let type_ns = crate::py_namespace!(vm, {
690            "__name__" => vm.ctx.new_str(self.type_name),
691            "__qualname__" => vm.ctx.new_str(self.type_qualname),
692            "__module__" => vm.ctx.new_str(self.type_module),
693        });
694        let msg_obj = match self.msg {
695            Some(m) => vm.ctx.new_str(m).into(),
696            None => vm.ctx.none(),
697        };
698        crate::py_namespace!(vm, {
699            "type" => type_ns,
700            "msg" => msg_obj,
701            "formatted" => vm.ctx.new_str(self.formatted),
702            "errdisplay" => vm.ctx.new_str(self.errdisplay),
703        })
704        .into()
705    }
706}
707
708/// `_PyCode_VerifyStateless`. `namespaces` holds the globals and the builtins
709/// the code's `LOAD_GLOBAL` names are resolved against; without them every such
710/// name stays unknown, which `_PyCode_CheckNoExternalState` lets through.
711fn verify_stateless(
712    code: &Py<PyCode>,
713    namespaces: Option<(&Py<PyDict>, &Py<PyDict>)>,
714    vm: &VirtualMachine,
715) -> PyResult<()> {
716    if !code.freevars.is_empty() {
717        return Err(vm.new_value_error("closures not supported"));
718    }
719    let Some((globals, builtins)) = namespaces else {
720        return Ok(());
721    };
722    // Having builtins to resolve against bumps numbuiltin by one, so a name
723    // that is neither a global nor a builtin is rejected along with the ones
724    // the globals do hold.
725    for name in global_names(code) {
726        if globals.contains_key(name, vm) || !builtins.contains_key(name, vm) {
727            return Err(vm.new_value_error("globals not supported"));
728        }
729    }
730    Ok(())
731}
732
733/// The `co_names` entries `identify_unbound_names` reaches through `LOAD_GLOBAL`.
734pub(crate) fn global_names(code: &Py<PyCode>) -> impl Iterator<Item = &'static PyStrInterned> + '_ {
735    walk_instructions(code).filter_map(|(_, instr, arg)| match instr {
736        Instruction::LoadGlobal { namei } => Some(code.names[(namei.get(arg) >> 1) as usize]),
737        _ => None,
738    })
739}
740
741/// The instruction stream with its inline caches skipped and its specialized
742/// and instrumented opcodes mapped back, yielding `(offset, instruction, arg)`.
743pub(crate) fn walk_instructions(
744    code: &Py<PyCode>,
745) -> impl Iterator<Item = (usize, Instruction, OpArg)> + '_ {
746    let units = &code.instructions;
747    let mut arg_state = OpArgState::default();
748    let mut offset = 0;
749    core::iter::from_fn(move || {
750        if offset >= units.len() {
751            return None;
752        }
753        let at = offset;
754        let op = units.read_op(at);
755        let arg = arg_state.extend(units.read_arg(at));
756        if !matches!(op, Instruction::ExtendedArg) {
757            arg_state.reset();
758        }
759        offset = at + 1 + op.cache_entries();
760        Some((at, op.deoptimize(), arg))
761    })
762}
763
764/// `_PyFunction_VerifyStateless`.
765fn verify_stateless_function(func: &Py<PyFunction>, vm: &VirtualMachine) -> PyResult<()> {
766    // `__globals__` is a dict by construction, so only the builtins are checked.
767    let builtins = func.builtins.downcast_ref::<PyDict>().ok_or_else(|| {
768        vm.new_type_error(format!(
769            "unsupported builtins {}",
770            render_repr(&func.builtins, vm)
771        ))
772    })?;
773    if func
774        .__defaults__()
775        .is_some_and(|d| !d.as_slice().is_empty())
776    {
777        return Err(vm.new_value_error("defaults not supported"));
778    }
779    if func.__kwdefaults__().is_some_and(|d| !d.is_empty()) {
780        return Err(vm.new_value_error("keyword defaults not supported"));
781    }
782    if func
783        .closure
784        .as_ref()
785        .is_some_and(|c| !c.as_slice().is_empty())
786    {
787        return Err(vm.new_value_error("closures not supported"));
788    }
789    verify_stateless(&func.code, Some((&func.globals, builtins)), vm)
790}
791
792/// `verify_script`: a script takes no arguments and returns only None.
793pub fn verify_script(code: &Py<PyCode>, vm: &VirtualMachine) -> PyResult<()> {
794    verify_stateless(code, None, vm)?;
795    if code.arg_count > 0
796        || code.posonlyarg_count > 0
797        || code.kwonlyarg_count > 0
798        || code.flags.contains(CodeFlags::VARARGS)
799        || code.flags.contains(CodeFlags::VARKEYWORDS)
800    {
801        return Err(vm.new_value_error("code with args not supported"));
802    }
803    if !code_returns_only_none(code) {
804        return Err(vm.new_value_error("code that returns a value is not a script"));
805    }
806    Ok(())
807}
808
809/// `_PyCode_CheckPureFunction`.
810fn is_pure_function(code: &Py<PyCode>) -> bool {
811    !code.flags.intersects(
812        CodeFlags::GENERATOR
813            | CodeFlags::COROUTINE
814            | CodeFlags::ITERABLE_COROUTINE
815            | CodeFlags::ASYNC_GENERATOR,
816    )
817}
818
819/// `_PyCode_ReturnsOnlyNone`. Here "value" means a non-None value, since a bare
820/// return is identical to returning None explicitly, as is a missing return
821/// statement at the end of the function.
822pub(crate) fn code_returns_only_none(code: &Py<PyCode>) -> bool {
823    if !is_pure_function(code) {
824        return false;
825    }
826    let units = &code.instructions;
827    let Some(last) = units.len().checked_sub(1) else {
828        return true;
829    };
830    // The last instruction either returns or raises. We can take advantage
831    // of that for a quick exit.
832    let final_op = units.read_op(last).deoptimize();
833
834    // Look up None in co_consts.
835    let Some(none_index) = code
836        .constants
837        .iter()
838        .position(|c| matches!(c.borrow_constant(), BorrowedConstant::None))
839    else {
840        // None wasn't there, which means there was no implicit return,
841        // "return", or "return None".
842
843        // That means there must be an explicit return (non-None), or it only
844        // raises.
845        if matches!(final_op, Instruction::ReturnValue) {
846            // It was an explicit return (non-None).
847            return false;
848        }
849        // It must end with a raise then. We still have to walk the bytecode
850        // to see if there's any explicit return (non-None).
851        return !walk_instructions(code).any(|(_, op, _)| matches!(op, Instruction::ReturnValue));
852    };
853    // Walk the bytecode, looking for RETURN_VALUE.
854    for (at, op, _) in walk_instructions(code) {
855        if !matches!(op, Instruction::ReturnValue) {
856            continue;
857        }
858        // Ignore it if it returns None.
859        if let Some(prev) = at.checked_sub(1)
860            && matches!(units.read_op(prev).deoptimize(), Instruction::LoadConst { .. })
861            // We don't worry about EXTENDED_ARG for now.
862            && usize::from(u8::from(units.read_arg(prev))) == none_index
863        {
864            continue;
865        }
866        return false;
867    }
868    true
869}
870
871/// Extract a script code object from source text / function / code object.
872pub fn script_code(obj: &PyObject, vm: &VirtualMachine) -> PyResult<PyRef<PyCode>> {
873    let code = if let Ok(code) = obj.to_owned().downcast::<PyCode>() {
874        code
875    } else if let Some(func) = obj.downcast_ref::<PyFunction>() {
876        (*func.code).to_owned()
877    } else {
878        #[allow(unused_variables)]
879        let src = source_as_string(obj, vm)?;
880        // Compile in the caller so SyntaxError is raised here.
881        #[cfg(feature = "rustpython-compiler")]
882        {
883            vm.compile(&src, crate::compiler::Mode::Exec, "<script>")
884                .map_err(|err| err.into_pyexception(vm, Some(src.as_str())))?
885        }
886        #[cfg(not(feature = "rustpython-compiler"))]
887        return Err(vm.new_type_error(
888            "can't compile a script to bytecode when the `codegen` feature of rustpython is disabled",
889        ));
890    };
891    verify_script(&code, vm)?;
892    Ok(code)
893}
894
895/// `_Py_SourceAsString` for the objects `_PyObject_SupportedAsScript` accepts:
896/// `str`, `bytes`, and other readable buffers.
897fn source_as_string(obj: &PyObject, vm: &VirtualMachine) -> PyResult<String> {
898    if let Some(s) = obj.downcast_ref::<PyStr>() {
899        return s
900            .as_wtf8()
901            .as_str()
902            .map(str::to_owned)
903            .map_err(|_| unsupported_script(vm, obj));
904    }
905    let buf = crate::function::ArgBytesLike::try_from_borrowed_object(vm, obj)
906        .map_err(|_| unsupported_script(vm, obj))?;
907    buf.with_ref(|bytes| {
908        core::str::from_utf8(bytes)
909            .map(str::to_owned)
910            .map_err(|_| unsupported_script(vm, obj))
911    })
912}
913
914fn unsupported_script(vm: &VirtualMachine, obj: &PyObject) -> PyBaseExceptionRef {
915    vm.new_type_error(format!("unsupported script {}", render_repr(obj, vm)))
916}
917
918pub fn apply_shared_ns(
919    ns: &crate::builtins::PyDictRef,
920    shared: &Py<crate::builtins::PyDict>,
921    vm: &VirtualMachine,
922) -> PyResult<()> {
923    for (key, value) in shared {
924        let name = utf8_key(&key, vm)?;
925        let shared = SharedValue::from_object(&value, Fallback::XidataOnly, vm)?;
926        ns.set_item(name, shared.into_object(vm)?, vm)?;
927    }
928    Ok(())
929}
930
931/// Run `f` in interpreter `id` with exclusive `__main__` occupancy.
932#[cfg(feature = "threading")]
933pub fn with_interpreter<F, R>(id: i64, caller: &VirtualMachine, f: F) -> PyResult<R>
934where
935    F: FnOnce(&VirtualMachine) -> PyResult<R>,
936{
937    let state = crate::vm::runtime::lookup_interpreter(id)
938        .ok_or_else(|| crate::stdlib::_interpreters::interpreter_not_found(caller, id))?;
939    if !state.ready.load(core::sync::atomic::Ordering::Acquire) {
940        return Err(crate::stdlib::_interpreters::interpreter_error(
941            caller,
942            format!("cannot exec interpreter {id} (not ready)"),
943        ));
944    }
945
946    if state
947        .running_main
948        .compare_exchange(
949            false,
950            true,
951            core::sync::atomic::Ordering::AcqRel,
952            core::sync::atomic::Ordering::Acquire,
953        )
954        .is_err()
955    {
956        return Err(crate::stdlib::_interpreters::interpreter_error(
957            caller,
958            "interpreter already running",
959        ));
960    }
961
962    struct RunningGuard<'a>(&'a crate::vm::PyGlobalState);
963    impl Drop for RunningGuard<'_> {
964        fn drop(&mut self) {
965            self.0
966                .running_main
967                .store(false, core::sync::atomic::Ordering::Release);
968        }
969    }
970    let _guard = RunningGuard(&state);
971
972    let tvm = crate::vm::runtime::owned_new_thread(id)
973        .ok_or_else(|| crate::stdlib::_interpreters::interpreter_not_found(caller, id))?;
974    tvm.run(f)
975}
976
977#[must_use]
978pub fn is_running(id: i64) -> bool {
979    crate::vm::runtime::lookup_interpreter(id)
980        .is_some_and(|s| s.is_main || s.running_main.load(core::sync::atomic::Ordering::Acquire))
981}