From 7c6510334cdedbf7a81ff2bf4ea0fa87499d6536 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" <69878+youknowone@users.noreply.github.com> Date: Sun, 13 Sep 2026 09:41:15 +0900 Subject: [PATCH] Keep import module names as PyUtf8Str (#8690) Add PyStr::as_utf8() for a non-erroring UTF-8 view. Thread PyUtf8Str through import resolution, module attribute errors, and IMPORT_FROM instead of converting names to owned String. Assisted-by: Grok:4.6 --- crates/vm/src/builtins/module.rs | 25 ++++------ crates/vm/src/builtins/str.rs | 30 ++++++++---- crates/vm/src/frame.rs | 20 +++++--- crates/vm/src/import.rs | 78 ++++++++++++++++++-------------- crates/vm/src/stdlib/sys.rs | 2 +- 5 files changed, 91 insertions(+), 64 deletions(-) diff --git a/crates/vm/src/builtins/module.rs b/crates/vm/src/builtins/module.rs index 8fa5259f7..9abbf459a 100644 --- a/crates/vm/src/builtins/module.rs +++ b/crates/vm/src/builtins/module.rs @@ -1,4 +1,4 @@ -use super::{PyDict, PyDictRef, PyStr, PyStrRef, PyType, PyTypeRef}; +use super::{PyDict, PyDictRef, PyStr, PyStrRef, PyType, PyTypeRef, PyUtf8Str}; use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, builtins::{PyStrInterned, pystr::AsPyStr}, @@ -161,17 +161,13 @@ impl Py { .get_item_opt(identifier!(vm, __name__), vm) .ok() .flatten(); - let mod_name_str = mod_name_obj.as_ref().and_then(|n| { - n.downcast_ref::() - .map(|s| s.to_string_lossy().into_owned()) - }); + let mod_name = mod_name_obj + .as_ref() + .and_then(|n| n.downcast_ref::()); // If __name__ is not set or not a string, use a simpler error message - let mod_display = match mod_name_str.as_deref() { - Some(s) => s, - None => { - return Err(vm.new_attribute_error(format!("module has no attribute '{name}'"))); - } + let Some(mod_display) = mod_name.map(|s| s.as_str()) else { + return Err(vm.new_attribute_error(format!("module has no attribute '{name}'"))); }; let spec = dict @@ -230,8 +226,7 @@ impl Py { } } else { // Check for uninitialized submodule - let submodule_initializing = - is_uninitialized_submodule(mod_name_str.as_deref(), name, vm); + let submodule_initializing = is_uninitialized_submodule(mod_name, name, vm); if submodule_initializing { Err(vm.new_attribute_error(format!( "cannot access submodule '{name}' of module '{mod_display}' \ @@ -465,7 +460,7 @@ pub(crate) fn init(context: &'static Context) { /// Check if {module_name}.{name} is an uninitialized submodule in sys.modules. fn is_uninitialized_submodule( - module_name: Option<&str>, + module_name: Option<&Py>, name: &Py, vm: &VirtualMachine, ) -> bool { @@ -477,8 +472,8 @@ fn is_uninitialized_submodule( return false; }; - let full_name = format!("{mod_name}.{name}"); - let Ok(sub_mod) = sys_modules.get_item(&full_name, vm) else { + let full_name = vm.ctx.new_utf8_str(format!("{}.{name}", mod_name.as_str())); + let Ok(sub_mod) = sys_modules.get_item(&*full_name, vm) else { return false; }; diff --git a/crates/vm/src/builtins/str.rs b/crates/vm/src/builtins/str.rs index 7e3328216..d9e08e3ec 100644 --- a/crates/vm/src/builtins/str.rs +++ b/crates/vm/src/builtins/str.rs @@ -613,20 +613,34 @@ impl PyStr { .map(|x| Self::from(unsafe { Wtf8Buf::from_bytes_unchecked(x) }).into_ref(&vm.ctx)) } + pub fn as_utf8(&self) -> Option<&PyUtf8Str> { + if self.is_utf8() { + // SAFETY: is_utf8() guarantees the PyUtf8Str invariant. + Some(unsafe { &*(self as *const Self as *const PyUtf8Str) }) + } else { + None + } + } + pub fn try_as_utf8<'a>(&'a self, vm: &VirtualMachine) -> PyResult<&'a PyUtf8Str> { - // Check if the string contains surrogates - self.ensure_valid_utf8(vm)?; - // If no surrogates, we can safely cast to PyStr - Ok(unsafe { &*(self as *const _ as *const PyUtf8Str) }) + self.as_utf8() + .ok_or_else(|| self.ensure_valid_utf8(vm).unwrap_err()) } } impl Py { + pub fn as_utf8(&self) -> Option<&Py> { + if self.is_utf8() { + // SAFETY: is_utf8() guarantees the PyUtf8Str invariant. + Some(unsafe { &*(self as *const Self as *const Py) }) + } else { + None + } + } + pub fn try_as_utf8<'a>(&'a self, vm: &VirtualMachine) -> PyResult<&'a Py> { - // Check if the string contains surrogates - self.ensure_valid_utf8(vm)?; - // If no surrogates, we can safely cast to PyStr - Ok(unsafe { &*(self as *const _ as *const Py) }) + self.as_utf8() + .ok_or_else(|| self.ensure_valid_utf8(vm).unwrap_err()) } } diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 61756b999..9600a92d9 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -7767,10 +7767,10 @@ impl ExecutingFrame<'_> { // fallback to importing '{module.__name__}.{name}' from sys.modules let fallback_module = (|| { let mod_name = module.get_attr(identifier!(vm, __name__), vm).ok()?; - let mod_name = mod_name.downcast_ref::()?; - let full_mod_name = format!("{mod_name}.{name}"); + let mod_name = mod_name.downcast_ref::()?; + let full_mod_name = vm.ctx.new_utf8_str(format!("{}.{name}", mod_name.as_str())); let sys_modules = vm.sys_module.get_attr("modules", vm).ok()?; - sys_modules.get_item(&full_mod_name, vm).ok() + sys_modules.get_item(&*full_mod_name, vm).ok() })(); if let Some(sub_module) = fallback_module { @@ -7783,10 +7783,10 @@ impl ExecutingFrame<'_> { // Get module name for the error message let mod_name_obj = module.get_attr(identifier!(vm, __name__), vm).ok(); - let mod_name_str = mod_name_obj + let mod_name = mod_name_obj .as_ref() - .and_then(|n| n.downcast_ref::().map(|s| s.as_str().to_owned())); - let module_name = mod_name_str.as_deref().unwrap_or(""); + .and_then(|n| n.downcast_ref::()); + let module_name = mod_name.map_or("", |s| s.as_str()); let spec = module .get_attr("__spec__", vm) @@ -7843,7 +7843,13 @@ impl ExecutingFrame<'_> { format!("cannot import name '{name}' from '{module_name}' (unknown location)") } }; - let err = vm.new_import_error(msg, vm.ctx.new_utf8_str(module_name)); + let err = vm.new_import_error( + msg, + match mod_name { + Some(s) => s.to_owned().into_wtf8(), + None => vm.ctx.new_utf8_str("").into_wtf8(), + }, + ); if let Some(ref path) = origin { let _ignore = err diff --git a/crates/vm/src/import.rs b/crates/vm/src/import.rs index eb07c76a2..38de9f044 100644 --- a/crates/vm/src/import.rs +++ b/crates/vm/src/import.rs @@ -162,7 +162,7 @@ pub fn import_source(vm: &VirtualMachine, module_name: &str, content: &str) -> P /// initializing by calling `_lock_unlock_module`. fn import_ensure_initialized( module: &PyObjectRef, - name: &str, + name: &Py, vm: &VirtualMachine, ) -> PyResult<()> { let initializing = match vm.get_attribute_opt(module.clone(), vm.ctx.intern_str("__spec__"))? { @@ -174,7 +174,7 @@ fn import_ensure_initialized( }; if initializing { let lock_unlock = vm.importlib.get_attr("_lock_unlock_module", vm)?; - lock_unlock.call((vm.ctx.new_utf8_str(name),), vm)?; + lock_unlock.call((name.to_owned(),), vm)?; } Ok(()) } @@ -383,11 +383,11 @@ pub(crate) fn import_module_level( return Err(vm.new_value_error("level must be >= 0")); } - let name_str = match name.to_str() { - Some(s) => s, + let name = match name.as_utf8() { + Some(name) => name, None => { - // Name contains surrogates. Like CPython, try sys.modules - // lookup with the Python string key directly. + // Name contains surrogates. Try sys.modules lookup with the + // Python string key directly. if level == 0 { let sys_modules = vm.sys_module.get_attr("modules", vm)?; return sys_modules.get_item(name, vm).map_err(|_| { @@ -421,12 +421,12 @@ pub(crate) fn import_module_level( vm.ctx.new_utf8_str(""), )); } - resolve_name(name_str, &package, level as usize, vm)? + resolve_name(name, &package, level as usize, vm)? } else { - if name_str.is_empty() { + if name.is_empty() { return Err(vm.new_value_error("Empty module name")); } - name_str.to_owned() + name.to_owned() }; // import_get_module + import_find_and_load @@ -438,8 +438,7 @@ pub(crate) fn import_module_level( } _ => { let find_and_load = vm.importlib.get_attr("_find_and_load", vm)?; - let abs_name_obj = vm.ctx.new_utf8_str(&*abs_name); - find_and_load.call((abs_name_obj, vm.import_func.clone()), vm)? + find_and_load.call((abs_name.clone(), vm.import_func.clone()), vm)? } }; @@ -464,15 +463,17 @@ pub(crate) fn import_module_level( } else { Ok(module) } - } else if level == 0 || !name_str.is_empty() { + } else if level == 0 || !name.is_empty() { + let name_str = name.as_str(); match name_str.find('.') { None => Ok(module), Some(dot) => { let to_return = if level == 0 { - name_str[..dot].to_owned() + vm.ctx.new_utf8_str(&name_str[..dot]) } else { let cut_off = name_str.len() - dot; - abs_name[..abs_name.len() - cut_off].to_owned() + let abs = abs_name.as_str(); + vm.ctx.new_utf8_str(&abs[..abs.len() - cut_off]) }; match sys_modules.get_item(&*to_return, vm) { Ok(m) => Ok(m), @@ -480,8 +481,7 @@ pub(crate) fn import_module_level( // For absolute imports (level 0), try importing the // parent. Matches _bootstrap.__import__ behavior. let find_and_load = vm.importlib.get_attr("_find_and_load", vm)?; - let to_return_obj = vm.ctx.new_utf8_str(&*to_return); - find_and_load.call((to_return_obj, vm.import_func.clone()), vm) + find_and_load.call((to_return, vm.import_func.clone()), vm) } Err(_) => { // For relative imports (level > 0), raise KeyError @@ -500,27 +500,38 @@ pub(crate) fn import_module_level( } /// resolve_name in import.c - resolve relative import name -fn resolve_name(name: &str, package: &str, level: usize, vm: &VirtualMachine) -> PyResult { +fn resolve_name( + name: &Py, + package: &Py, + level: usize, + vm: &VirtualMachine, +) -> PyResult { // Python: bits = package.rsplit('.', level - 1) // Rust: rsplitn(level, '.') gives maxsplit=level-1 - let parts: Vec<&str> = package.rsplitn(level, '.').collect(); + let package_str = package.as_str(); + let parts: Vec<&str> = package_str.rsplitn(level, '.').collect(); if parts.len() < level { return Err(vm.new_import_error( "attempted relative import beyond top-level package", - vm.ctx.new_utf8_str(name), + name.to_owned(), )); } // rsplitn returns parts right-to-left, so last() is the leftmost (base) - let base = parts.last().unwrap(); - if name.is_empty() { - Ok(base.to_string()) + let base = *parts.last().unwrap(); + let abs_name = if name.is_empty() { + if base.len() == package_str.len() { + package.to_owned() + } else { + vm.ctx.new_utf8_str(base) + } } else { - Ok(format!("{base}.{name}")) - } + vm.ctx.new_utf8_str(format!("{base}.{}", name.as_str())) + }; + Ok(abs_name) } /// _calc___package__ - calculate package from globals for relative imports -fn calc_package(globals: Option<&PyObjectRef>, vm: &VirtualMachine) -> PyResult { +fn calc_package(globals: Option<&PyObjectRef>, vm: &VirtualMachine) -> PyResult { let globals = globals.ok_or_else(|| { vm.new_import_error( "attempted relative import with no known parent package", @@ -570,7 +581,7 @@ fn calc_package(globals: Option<&PyObjectRef>, vm: &VirtualMachine) -> PyResult< ); } } - return Ok(pkg_str.as_str().to_owned()); + return Ok(pkg_str); } else if let Some(ref spec) = spec && !vm.is_none(spec) && let Ok(parent) = spec.get_attr("parent", vm) @@ -579,7 +590,7 @@ fn calc_package(globals: Option<&PyObjectRef>, vm: &VirtualMachine) -> PyResult< let parent_str: PyUtf8StrRef = parent .downcast() .map_err(|_| vm.new_type_error("package set to non-string"))?; - return Ok(parent_str.as_str().to_owned()); + return Ok(parent_str); } // Fall back to __name__ and __path__ @@ -605,14 +616,15 @@ fn calc_package(globals: Option<&PyObjectRef>, vm: &VirtualMachine) -> PyResult< let mod_name_str: PyUtf8StrRef = mod_name .downcast() .map_err(|_| vm.new_type_error("__name__ must be a string"))?; - let mut package = mod_name_str.as_str().to_owned(); // If not a package (no __path__), strip last component. // Uses rpartition('.')[0] semantics: returns empty string when no dot. if globals.get_item("__path__", vm).is_err() { - package = match package.rfind('.') { - Some(dot) => package[..dot].to_owned(), - None => String::new(), - }; + let s = mod_name_str.as_str(); + Ok(match s.rfind('.') { + Some(dot) => vm.ctx.new_utf8_str(&s[..dot]), + None => vm.ctx.new_utf8_str(""), + }) + } else { + Ok(mod_name_str) } - Ok(package) } diff --git a/crates/vm/src/stdlib/sys.rs b/crates/vm/src/stdlib/sys.rs index 46784509c..166f9d829 100644 --- a/crates/vm/src/stdlib/sys.rs +++ b/crates/vm/src/stdlib/sys.rs @@ -309,7 +309,7 @@ pub mod sys { vm.ctx.new_tuple( module_names .into_iter() - .map(|n| vm.ctx.new_str(n).into()) + .map(|n| vm.ctx.new_utf8_str(n).into()) .collect(), ) }