1#[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
28pub(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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum Fallback {
67 XidataOnly = 0,
69 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
89pub 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#[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 Code(Vec<u8>),
141 Function(Vec<u8>),
143 Pickled(PickledData),
145}
146
147#[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 match pickle_dumps(obj, vm) {
169 Ok(data) => Ok(Self::Pickled(data)),
170 Err(_) => Err(exc),
172 }
173 }
174 }
175 }
176
177 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 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 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 if cls.is(vm.ctx.types.memoryview_type) && memoryview_registered() {
241 return Some(Self::from_buffer_object(obj, vm));
242 }
243 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 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 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 Err(_) => Err(exc),
283 },
284 }
285 }
286
287 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 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
364fn 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 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 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 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 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 let namespace = run_path.call((mainfile, vm.ctx.none(), "<fake __main__>"), vm)?;
506 loaded.dict().merge_object(namespace, vm)?;
507 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
515fn 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
526fn 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
547pub 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
567pub 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
583fn 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 if formatted.ends_with('\n') {
610 formatted.pop();
611 }
612 Ok(formatted)
613}
614
615fn 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
629pub 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 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
708fn 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 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
733pub(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
741pub(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
764fn verify_stateless_function(func: &Py<PyFunction>, vm: &VirtualMachine) -> PyResult<()> {
766 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
792pub 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
809fn 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
819pub(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 let final_op = units.read_op(last).deoptimize();
833
834 let Some(none_index) = code
836 .constants
837 .iter()
838 .position(|c| matches!(c.borrow_constant(), BorrowedConstant::None))
839 else {
840 if matches!(final_op, Instruction::ReturnValue) {
846 return false;
848 }
849 return !walk_instructions(code).any(|(_, op, _)| matches!(op, Instruction::ReturnValue));
852 };
853 for (at, op, _) in walk_instructions(code) {
855 if !matches!(op, Instruction::ReturnValue) {
856 continue;
857 }
858 if let Some(prev) = at.checked_sub(1)
860 && matches!(units.read_op(prev).deoptimize(), Instruction::LoadConst { .. })
861 && usize::from(u8::from(units.read_arg(prev))) == none_index
863 {
864 continue;
865 }
866 return false;
867 }
868 true
869}
870
871pub 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 #[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
895fn 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#[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}