Skip to content

Commit e1fe46f

Browse files
committed
Implement maybe_pyc_file and .pyc file execution
1 parent 72a90da commit e1fe46f

File tree

6 files changed

+69
-30
lines changed

6 files changed

+69
-30
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/import.rs

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ use crate::{
55
builtins::{PyCode, list, traceback::PyTraceback},
66
exceptions::types::PyBaseException,
77
scope::Scope,
8-
version::get_git_revision,
98
vm::{VirtualMachine, thread},
109
};
1110

@@ -44,16 +43,6 @@ pub(crate) fn init_importlib_package(vm: &VirtualMachine, importlib: PyObjectRef
4443

4544
let install_external = importlib.get_attr("_install_external_importers", vm)?;
4645
install_external.call((), vm)?;
47-
// Set pyc magic number to commit hash. Should be changed when bytecode will be more stable.
48-
let importlib_external = vm.import("_frozen_importlib_external", 0)?;
49-
let mut magic = get_git_revision().into_bytes();
50-
magic.truncate(4);
51-
if magic.len() != 4 {
52-
// os_random is expensive, but this is only ever called once
53-
magic = rustpython_common::rand::os_random::<4>().to_vec();
54-
}
55-
let magic: PyObjectRef = vm.ctx.new_bytes(magic).into();
56-
importlib_external.set_attr("MAGIC_NUMBER", magic, vm)?;
5746
let zipimport_res = (|| -> PyResult<()> {
5847
let zipimport = vm.import("zipimport", 0)?;
5948
let zipimporter = zipimport.get_attr("zipimporter", vm)?;

crates/vm/src/stdlib/imp.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ mod _imp {
8585
PyObjectRef, PyRef, PyResult, VirtualMachine,
8686
builtins::{PyBytesRef, PyCode, PyMemoryView, PyModule, PyStrRef},
8787
function::OptionalArg,
88-
import,
88+
import, version,
8989
};
9090

9191
#[pyattr]
@@ -94,6 +94,9 @@ mod _imp {
9494
.new_str(vm.state.config.settings.check_hash_pycs_mode.to_string())
9595
}
9696

97+
#[pyattr(name = "pyc_magic_number_token")]
98+
use version::PYC_MAGIC_NUMBER_TOKEN;
99+
97100
#[pyfunction]
98101
const fn extension_suffixes() -> PyResult<Vec<PyObjectRef>> {
99102
Ok(Vec::new())

crates/vm/src/version.rs

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

127127
format!("{date} {time}")
128128
}
129+
130+
// Must be aligned to Lib/importlib/_bootstrap_external.py
131+
pub const PYC_MAGIC_NUMBER: u16 = 3531;
132+
133+
// CPython format: magic_number | ('\r' << 16) | ('\n' << 24)
134+
// This protects against text-mode file reads
135+
pub const PYC_MAGIC_NUMBER_TOKEN: u32 =
136+
(PYC_MAGIC_NUMBER as u32) | ((b'\r' as u32) << 16) | ((b'\n' as u32) << 24);
137+
138+
/// Magic number as little-endian bytes for .pyc files
139+
pub const PYC_MAGIC_NUMBER_BYTES: [u8; 4] = PYC_MAGIC_NUMBER_TOKEN.to_le_bytes();

crates/vm/src/vm/compile.rs

Lines changed: 54 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,57 @@ 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+
maybe_pyc_file_with_magic(path, &crate::version::PYC_MAGIC_NUMBER_BYTES).unwrap_or(false)
162+
}
163+
164+
fn maybe_pyc_file_with_magic(path: &str, magic_number: &[u8]) -> std::io::Result<bool> {
165+
// For non-.pyc extension, check magic number
166+
let path_obj = std::path::Path::new(path);
167+
if !path_obj.is_file() {
168+
return Ok(false);
169+
}
170+
171+
let mut file = std::fs::File::open(path)?;
172+
let mut buf = [0u8; 2];
173+
174+
use std::io::Read;
175+
if file.read(&mut buf)? != 2 || magic_number.len() < 2 {
176+
return Ok(false);
177+
}
178+
179+
// Read only two bytes of the magic. If the file was opened in
180+
// text mode, the bytes 3 and 4 of the magic (\r\n) might not
181+
// be read as they are on disk.
182+
// See CPython's pythonrun.c:maybe_pyc_file for reference.
183+
Ok(buf == magic_number[..2])
184+
}
185+
138186
fn get_importer(path: &str, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
139187
let path_importer_cache = vm.sys_module.get_attr("path_importer_cache", vm)?;
140188
let path_importer_cache = PyDictRef::try_from_object(vm, path_importer_cache)?;

0 commit comments

Comments
 (0)