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

@@ -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<PyModule> {
             .get_item_opt(identifier!(vm, __name__), vm)
             .ok()
             .flatten();
-        let mod_name_str = mod_name_obj.as_ref().and_then(|n| {
-            n.downcast_ref::<PyStr>()
-                .map(|s| s.to_string_lossy().into_owned())
-        });
+        let mod_name = mod_name_obj
+            .as_ref()
+            .and_then(|n| n.downcast_ref::<PyUtf8Str>());
 
         // 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<PyModule> {
                 }
             } 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<PyUtf8Str>>,
     name: &Py<PyStr>,
     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;
     };
 

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<PyStr> {
+    pub fn as_utf8(&self) -> Option<&Py<PyUtf8Str>> {
+        if self.is_utf8() {
+            // SAFETY: is_utf8() guarantees the PyUtf8Str invariant.
+            Some(unsafe { &*(self as *const Self as *const Py<PyUtf8Str>) })
+        } else {
+            None
+        }
+    }
+
     pub fn try_as_utf8<'a>(&'a self, vm: &VirtualMachine) -> PyResult<&'a Py<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 Py<PyUtf8Str>) })
+        self.as_utf8()
+            .ok_or_else(|| self.ensure_valid_utf8(vm).unwrap_err())
     }
 }
 

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::<PyStr>()?;
-            let full_mod_name = format!("{mod_name}.{name}");
+            let mod_name = mod_name.downcast_ref::<PyUtf8Str>()?;
+            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::<PyUtf8Str>().map(|s| s.as_str().to_owned()));
-        let module_name = mod_name_str.as_deref().unwrap_or("<unknown module name>");
+            .and_then(|n| n.downcast_ref::<PyUtf8Str>());
+        let module_name = mod_name.map_or("<unknown module name>", |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("<unknown module name>").into_wtf8(),
+            },
+        );
 
         if let Some(ref path) = origin {
             let _ignore = err

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<PyUtf8Str>,
     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<String> {
+fn resolve_name(
+    name: &Py<PyUtf8Str>,
+    package: &Py<PyUtf8Str>,
+    level: usize,
+    vm: &VirtualMachine,
+) -> PyResult<PyUtf8StrRef> {
     // 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<String> {
+fn calc_package(globals: Option<&PyObjectRef>, vm: &VirtualMachine) -> PyResult<PyUtf8StrRef> {
     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)
 }

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(),
         )
     }