Skip to content

Commit 5e90531

Browse files
youknowoneyouknowone
committed
vm.run_pyc_bytes
1 parent 7e33efa commit 5e90531

File tree

2 files changed

+95
-48
lines changed

2 files changed

+95
-48
lines changed

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

Lines changed: 84 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@ mod vm_ops;
1919
use crate::{
2020
AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
2121
builtins::{
22-
PyBaseExceptionRef, PyDictRef, PyInt, PyList, PyModule, PyStr, PyStrInterned, PyStrRef,
23-
PyTypeRef, code::PyCode, pystr::AsPyStr, tuple::PyTuple,
22+
PyBaseExceptionRef, PyDict, PyDictRef, PyInt, PyList, PyModule, PyStr, PyStrInterned,
23+
PyStrRef, PyTypeRef, code::PyCode, pystr::AsPyStr, tuple::PyTuple,
2424
},
2525
codecs::CodecsRegistry,
2626
common::{hash::HashSecret, lock::PyMutex, rc::PyRc},
@@ -460,6 +460,42 @@ impl VirtualMachine {
460460
self.signal_rx = Some(signal_rx);
461461
}
462462

463+
/// Execute Python bytecode (`.pyc`) from an in-memory buffer.
464+
///
465+
/// When the RustPython CLI is available, `.pyc` files are normally executed by
466+
/// invoking `rustpython <input>.pyc`. This method provides an alternative for
467+
/// environments where the binary is unavailable or file I/O is restricted
468+
/// (e.g. WASM).
469+
///
470+
/// ## Preparing a `.pyc` file
471+
///
472+
/// First, compile a Python source file into bytecode:
473+
///
474+
/// ```sh
475+
/// # Generate a .pyc file
476+
/// $ rustpython -m py_compile <input>.py
477+
/// ```
478+
///
479+
/// ## Running the bytecode
480+
///
481+
/// Load the resulting `.pyc` file into memory and execute it using the VM:
482+
///
483+
/// ```no_run
484+
/// use rustpython_vm::Interpreter;
485+
/// Interpreter::without_stdlib(Default::default()).enter(|vm| {
486+
/// let bytes = std::fs::read("__pycache__/<input>.rustpython-313.pyc").unwrap();
487+
/// let main_scope = vm.new_scope_with_main().unwrap();
488+
/// vm.run_pyc_bytes(&bytes, main_scope);
489+
/// });
490+
/// ```
491+
pub fn run_pyc_bytes(&self, pyc_bytes: &[u8], scope: Scope) -> PyResult<()> {
492+
let code = PyCode::from_pyc(pyc_bytes, Some("<pyc_bytes>"), None, None, self)?;
493+
self.with_simple_run("<source>", |_module_dict| {
494+
self.run_code_obj(code, scope)?;
495+
Ok(())
496+
})
497+
}
498+
463499
pub fn run_code_obj(&self, code: PyRef<PyCode>, scope: Scope) -> PyResult {
464500
use crate::builtins::PyFunction;
465501

@@ -500,6 +536,52 @@ impl VirtualMachine {
500536
}
501537
}
502538

539+
/// Run `run` with main scope.
540+
fn with_simple_run(
541+
&self,
542+
path: &str,
543+
run: impl FnOnce(&Py<PyDict>) -> PyResult<()>,
544+
) -> PyResult<()> {
545+
let sys_modules = self.sys_module.get_attr(identifier!(self, modules), self)?;
546+
let main_module = sys_modules.get_item(identifier!(self, __main__), self)?;
547+
let module_dict = main_module.dict().expect("main module must have __dict__");
548+
549+
// Track whether we set __file__ (for cleanup)
550+
let set_file_name = !module_dict.contains_key(identifier!(self, __file__), self);
551+
if set_file_name {
552+
module_dict.set_item(
553+
identifier!(self, __file__),
554+
self.ctx.new_str(path).into(),
555+
self,
556+
)?;
557+
module_dict.set_item(identifier!(self, __cached__), self.ctx.none(), self)?;
558+
}
559+
560+
let result = run(&module_dict);
561+
562+
self.flush_io();
563+
564+
// Cleanup __file__ and __cached__ after execution
565+
if set_file_name {
566+
let _ = module_dict.del_item(identifier!(self, __file__), self);
567+
let _ = module_dict.del_item(identifier!(self, __cached__), self);
568+
}
569+
570+
result
571+
}
572+
573+
/// flush_io
574+
///
575+
/// Flush stdout and stderr. Errors are silently ignored.
576+
fn flush_io(&self) {
577+
if let Ok(stdout) = self.sys_module.get_attr("stdout", self) {
578+
let _ = self.call_method(&stdout, identifier!(self, flush).as_str(), ());
579+
}
580+
if let Ok(stderr) = self.sys_module.get_attr("stderr", self) {
581+
let _ = self.call_method(&stderr, identifier!(self, flush).as_str(), ());
582+
}
583+
}
584+
503585
pub fn current_recursion_depth(&self) -> usize {
504586
self.recursion_depth.get()
505587
}

‎crates/vm/src/vm/python_run.rs‎

Lines changed: 11 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
//! Python code execution functions.
22
33
use crate::{
4-
PyResult, VirtualMachine,
5-
builtins::{PyCode, PyDictRef},
4+
Py, PyResult, VirtualMachine,
5+
builtins::{PyCode, PyDict},
66
compiler::{self},
77
scope::Scope,
88
};
@@ -25,37 +25,14 @@ impl VirtualMachine {
2525
/// Execute a Python file with __main__ module setup.
2626
/// Sets __file__ and __cached__ before execution, removes them after.
2727
fn run_simple_file(&self, scope: Scope, path: &str) -> PyResult<()> {
28-
let sys_modules = self.sys_module.get_attr(identifier!(self, modules), self)?;
29-
let main_module = sys_modules.get_item(identifier!(self, __main__), self)?;
30-
let module_dict = main_module.dict().expect("main module must have __dict__");
31-
32-
// Track whether we set __file__ (for cleanup)
33-
let set_file_name = !module_dict.contains_key(identifier!(self, __file__), self);
34-
if set_file_name {
35-
module_dict.set_item(
36-
identifier!(self, __file__),
37-
self.ctx.new_str(path).into(),
38-
self,
39-
)?;
40-
module_dict.set_item(identifier!(self, __cached__), self.ctx.none(), self)?;
41-
}
42-
43-
let result = self.run_simple_file_inner(&module_dict, scope, path);
44-
45-
self.flush_io();
46-
47-
// Cleanup __file__ and __cached__ after execution
48-
if set_file_name {
49-
let _ = module_dict.del_item(identifier!(self, __file__), self);
50-
let _ = module_dict.del_item(identifier!(self, __cached__), self);
51-
}
52-
53-
result
28+
self.with_simple_run(path, |module_dict| {
29+
self.run_simple_file_inner(module_dict, scope, path)
30+
})
5431
}
5532

5633
fn run_simple_file_inner(
5734
&self,
58-
module_dict: &PyDictRef,
35+
module_dict: &Py<PyDict>,
5936
scope: Scope,
6037
path: &str,
6138
) -> PyResult<()> {
@@ -123,22 +100,10 @@ impl VirtualMachine {
123100
.map_err(|err| self.new_syntax_error(&err, Some(source)))?;
124101
self.run_code_obj(code_obj, scope)
125102
}
126-
127-
/// flush_io
128-
///
129-
/// Flush stdout and stderr. Errors are silently ignored.
130-
fn flush_io(&self) {
131-
if let Ok(stdout) = self.sys_module.get_attr("stdout", self) {
132-
let _ = self.call_method(&stdout, identifier!(self, flush).as_str(), ());
133-
}
134-
if let Ok(stderr) = self.sys_module.get_attr("stderr", self) {
135-
let _ = self.call_method(&stderr, identifier!(self, flush).as_str(), ());
136-
}
137-
}
138103
}
139104

140105
fn set_main_loader(
141-
module_dict: &PyDictRef,
106+
module_dict: &Py<PyDict>,
142107
filename: &str,
143108
loader_name: &str,
144109
vm: &VirtualMachine,
@@ -162,10 +127,10 @@ fn maybe_pyc_file(path: &str) -> bool {
162127
if path.ends_with(".pyc") {
163128
return true;
164129
}
165-
maybe_pyc_file_with_magic(path, &crate::version::PYC_MAGIC_NUMBER_BYTES).unwrap_or(false)
130+
maybe_pyc_file_with_magic(path).unwrap_or(false)
166131
}
167132

168-
fn maybe_pyc_file_with_magic(path: &str, magic_number: &[u8]) -> std::io::Result<bool> {
133+
fn maybe_pyc_file_with_magic(path: &str) -> std::io::Result<bool> {
169134
let path_obj = std::path::Path::new(path);
170135
if !path_obj.is_file() {
171136
return Ok(false);
@@ -175,12 +140,12 @@ fn maybe_pyc_file_with_magic(path: &str, magic_number: &[u8]) -> std::io::Result
175140
let mut buf = [0u8; 2];
176141

177142
use std::io::Read;
178-
if file.read(&mut buf)? != 2 || magic_number.len() < 2 {
143+
if file.read(&mut buf)? != 2 {
179144
return Ok(false);
180145
}
181146

182147
// Read only two bytes of the magic. If the file was opened in
183148
// text mode, the bytes 3 and 4 of the magic (\r\n) might not
184149
// be read as they are on disk.
185-
Ok(buf == magic_number[..2])
150+
Ok(crate::import::check_pyc_magic_number_bytes(&buf))
186151
}

0 commit comments

Comments
 (0)