Skip to content

Commit c979a0b

Browse files
youknowoneyouknowone
committed
Implement maybe_pyc_file and .pyc file execution
1 parent c9bf8df commit c979a0b

File tree

4 files changed

+62
-18
lines changed

4 files changed

+62
-18
lines changed

‎Lib/test/test_cmd_line_script.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -247,8 +247,6 @@ def test_script_abspath(self):
247247
script_dir, None,
248248
importlib.machinery.SourceFileLoader)
249249

250-
# TODO: RUSTPYTHON
251-
@unittest.expectedFailure
252250
def test_script_compiled(self):
253251
with os_helper.temp_dir() as script_dir:
254252
script_name = _make_test_script(script_dir, 'script')

‎Lib/test/test_compileall.py‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -325,8 +325,6 @@ def _test_ddir_only(self, *, ddir, parallel=True):
325325
self.assertEqual(mod_code_obj.co_filename, expected_in)
326326
self.assertIn(f'"{expected_in}"', os.fsdecode(err))
327327

328-
# TODO: RUSTPYTHON
329-
@unittest.expectedFailure
330328
def test_ddir_only_one_worker(self):
331329
"""Recursive compile_dir ddir= contains package paths; bpo39769."""
332330
return self._test_ddir_only(ddir="<a prefix>", parallel=False)
@@ -336,8 +334,6 @@ def test_ddir_multiple_workers(self):
336334
"""Recursive compile_dir ddir= contains package paths; bpo39769."""
337335
return self._test_ddir_only(ddir="<a prefix>", parallel=True)
338336

339-
# TODO: RUSTPYTHON
340-
@unittest.expectedFailure
341337
def test_ddir_empty_only_one_worker(self):
342338
"""Recursive compile_dir ddir='' contains package paths; bpo39769."""
343339
return self._test_ddir_only(ddir="", parallel=False)
@@ -347,8 +343,6 @@ def test_ddir_empty_multiple_workers(self):
347343
"""Recursive compile_dir ddir='' contains package paths; bpo39769."""
348344
return self._test_ddir_only(ddir="", parallel=True)
349345

350-
# TODO: RUSTPYTHON
351-
@unittest.expectedFailure
352346
def test_strip_only(self):
353347
fullpath = ["test", "build", "real", "path"]
354348
path = os.path.join(self.directory, *fullpath)
@@ -408,8 +402,6 @@ def test_prepend_only(self):
408402
str(err, encoding=sys.getdefaultencoding())
409403
)
410404

411-
# TODO: RUSTPYTHON
412-
@unittest.expectedFailure
413405
def test_strip_and_prepend(self):
414406
fullpath = ["test", "build", "real", "path"]
415407
path = os.path.join(self.directory, *fullpath)
@@ -887,8 +879,6 @@ def test_workers_available_cores(self, compile_dir):
887879
self.assertTrue(compile_dir.called)
888880
self.assertEqual(compile_dir.call_args[-1]['workers'], 0)
889881

890-
# TODO: RUSTPYTHON
891-
@unittest.expectedFailure
892882
def test_strip_and_prepend(self):
893883
fullpath = ["test", "build", "real", "path"]
894884
path = os.path.join(self.directory, *fullpath)

‎crates/vm/src/version.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,3 +126,13 @@ pub fn get_git_datetime() -> String {
126126

127127
format!("{date} {time}")
128128
}
129+
130+
/// Returns the magic number for .pyc files (4 bytes from git revision)
131+
pub fn get_magic_number() -> Vec<u8> {
132+
let mut magic = get_git_revision().into_bytes();
133+
magic.truncate(4);
134+
if magic.len() != 4 {
135+
magic = rustpython_common::rand::os_random::<4>().to_vec();
136+
}
137+
magic
138+
}

‎crates/vm/src/vm/compile.rs‎

Lines changed: 52 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -75,13 +75,20 @@ impl VirtualMachine {
7575
// Consider to use enum to distinguish `path`
7676
// https://github.com/RustPython/RustPython/pull/6276#discussion_r2529849479
7777

78-
// TODO: check .pyc here
79-
let pyc = false;
78+
let pyc = maybe_pyc_file(path);
8079
if pyc {
81-
todo!("running pyc is not implemented yet");
80+
// pyc file execution
81+
set_main_loader(&module_dict, path, "SourcelessFileLoader", self)?;
82+
let loader = module_dict.get_item("__loader__", self)?;
83+
let get_code = loader.get_attr("get_code", self)?;
84+
let code_obj = get_code.call((identifier!(self, __main__).to_owned(),), self)?;
85+
let code = code_obj
86+
.downcast::<PyCode>()
87+
.map_err(|_| self.new_runtime_error("Bad code object in .pyc file".to_owned()))?;
88+
self.run_code_obj(code, scope)?;
8289
} else {
8390
if path != "<stdin>" {
84-
set_main_loader(&module_dict, path, self)?;
91+
set_main_loader(&module_dict, path, "SourceFileLoader", self)?;
8592
}
8693
// TODO: replace to something equivalent to py_run_file
8794
match std::fs::read_to_string(path) {
@@ -125,16 +132,55 @@ impl VirtualMachine {
125132
}
126133
}
127134

128-
fn set_main_loader(module_dict: &PyDictRef, filename: &str, vm: &VirtualMachine) -> PyResult<()> {
135+
fn set_main_loader(
136+
module_dict: &PyDictRef,
137+
filename: &str,
138+
loader_name: &str,
139+
vm: &VirtualMachine,
140+
) -> PyResult<()> {
129141
vm.import("importlib.machinery", 0)?;
130142
let sys_modules = vm.sys_module.get_attr(identifier!(vm, modules), vm)?;
131143
let machinery = sys_modules.get_item("importlib.machinery", vm)?;
132-
let loader_class = machinery.get_attr("SourceFileLoader", vm)?;
144+
let loader_name = vm.ctx.new_str(loader_name);
145+
let loader_class = machinery.get_attr(&loader_name, vm)?;
133146
let loader = loader_class.call((identifier!(vm, __main__).to_owned(), filename), vm)?;
134147
module_dict.set_item("__loader__", loader, vm)?;
135148
Ok(())
136149
}
137150

151+
/// Check whether a file is maybe a pyc file.
152+
///
153+
/// Detection is performed by:
154+
/// 1. Checking if the filename ends with ".pyc"
155+
/// 2. If not, reading the first 2 bytes and comparing with the magic number
156+
fn maybe_pyc_file(path: &str) -> bool {
157+
// 1. Check if filename ends with ".pyc"
158+
if path.ends_with(".pyc") {
159+
return true;
160+
}
161+
let magic_number = crate::version::get_magic_number();
162+
maybe_pyc_file_with_magic(path, &magic_number).unwrap_or(false)
163+
}
164+
165+
fn maybe_pyc_file_with_magic(path: &str, magic_number: &[u8]) -> std::io::Result<bool> {
166+
// 2. For non-.pyc extension, check magic number
167+
let path_obj = std::path::Path::new(path);
168+
if !path_obj.is_file() {
169+
return Ok(false);
170+
}
171+
172+
let mut file = std::fs::File::open(path)?;
173+
let mut buf = [0u8; 2];
174+
175+
use std::io::Read;
176+
if file.read(&mut buf)? != 2 || magic_number.len() < 2 {
177+
return Ok(false);
178+
}
179+
180+
// Compare with first 2 bytes of magic number (half magic)
181+
Ok(buf == magic_number[..2])
182+
}
183+
138184
fn get_importer(path: &str, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
139185
let path_importer_cache = vm.sys_module.get_attr("path_importer_cache", vm)?;
140186
let path_importer_cache = PyDictRef::try_from_object(vm, path_importer_cache)?;

0 commit comments

Comments
 (0)