Merge pull request RustPython#4559 from youknowone/importlib-requirements

youknowone · web-flow · commit f2f8ecac733d · 2023-02-25T22:59:14.000+09:00
Add (fake) frozenmodule __origname__ and _imp.find_frozen
diff --git a/.cspell.json b/.cspell.json
@@ -65,6 +65,7 @@
     "PYTHONWARNINGS",
     "basicsize",
     "itemsize",
+    "origname",
     "getattro",
     "setattro",
     "iternext",
diff --git a/vm/build.rs b/vm/build.rs
@@ -2,11 +2,16 @@ use itertools::Itertools;
 use std::{env, io::prelude::*, path::PathBuf, process::Command};
 
 fn main() {
-    #[cfg(feature = "freeze-stdlib")]
-    for entry in glob::glob("Lib/*/*.py").expect("Lib/ exists?").flatten() {
+    let frozen_libs = if cfg!(feature = "freeze-stdlib") {
+        "Lib/*/*.py"
+    } else {
+        "Lib/python_builtins/*.py"
+    };
+    for entry in glob::glob(frozen_libs).expect("Lib/ exists?").flatten() {
         let display = entry.display();
         println!("cargo:rerun-if-changed={display}");
     }
+    println!("cargo:rerun-if-changed=../Lib/importlib/_bootstrap.py");
 
     println!("cargo:rustc-env=RUSTPYTHON_GIT_HASH={}", git_hash());
     println!(
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -21,11 +21,11 @@ pub(crate) fn init_importlib_base(vm: &mut VirtualMachine) -> PyResult<PyObjectR
     import_builtin(vm, "_weakref")?;
 
     let importlib = thread::enter_vm(vm, || {
-        let importlib = import_frozen(vm, "_frozen_importlib")?;
-        let impmod = import_builtin(vm, "_imp")?;
-        let install = importlib.get_attr("_install", vm)?;
-        vm.invoke(&install, (vm.sys_module.clone(), impmod))?;
-        Ok(importlib)
+        let bootstrap = import_frozen(vm, "_frozen_importlib")?;
+        let install = bootstrap.get_attr("_install", vm)?;
+        let imp = import_builtin(vm, "_imp")?;
+        vm.invoke(&install, (vm.sys_module.clone(), imp))?;
+        Ok(bootstrap)
     })?;
     vm.import_func = importlib.get_attr(identifier!(vm, __import__).to_owned(), vm)?;
     Ok(importlib)
@@ -81,7 +81,12 @@ pub fn make_frozen(vm: &VirtualMachine, name: &str) -> PyResult<PyRef<PyCode>> {
 }
 
 pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
-    make_frozen(vm, module_name).and_then(|frozen| import_codeobj(vm, module_name, frozen, false))
+    make_frozen(vm, module_name).and_then(|frozen| {
+        let module = import_codeobj(vm, module_name, frozen, false)?;
+        // TODO: give a correct origname here
+        module.set_attr("__origname__", vm.ctx.new_str(module_name.to_owned()), vm)?;
+        Ok(module)
+    })
 }
 
 pub fn import_builtin(vm: &VirtualMachine, module_name: &str) -> PyResult {
diff --git a/vm/src/stdlib/imp.rs b/vm/src/stdlib/imp.rs
@@ -1,4 +1,4 @@
-use crate::{PyObjectRef, VirtualMachine};
+use crate::{builtins::PyBaseExceptionRef, bytecode::FrozenModule, PyObjectRef, VirtualMachine};
 
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let module = _imp::make_module(vm);
@@ -48,10 +48,40 @@ mod lock {
     }
 }
 
+#[allow(dead_code)]
+enum FrozenError {
+    BadName,  // The given module name wasn't valid.
+    NotFound, // It wasn't in PyImport_FrozenModules.
+    Disabled, // -X frozen_modules=off (and not essential)
+    Excluded, // The PyImport_FrozenModules entry has NULL "code"
+    //        (module is present but marked as unimportable, stops search).
+    Invalid, // The PyImport_FrozenModules entry is bogus
+             //          (eg. does not contain executable code).
+}
+
+impl FrozenError {
+    fn to_pyexception(&self, mod_name: &str, vm: &VirtualMachine) -> PyBaseExceptionRef {
+        use FrozenError::*;
+        let msg = match self {
+            BadName | NotFound => format!("No such frozen object named {mod_name}"),
+            Disabled => format!("Frozen modules are disabled and the frozen object named {mod_name} is not essential"),
+            Excluded => format!("Excluded frozen object named {mod_name}"),
+            Invalid => format!("Frozen object named {mod_name} is invalid"),
+        };
+        vm.new_import_error(msg, mod_name)
+    }
+}
+
+// find_frozen in frozen.c
+fn find_frozen<'a>(name: &str, vm: &'a VirtualMachine) -> Result<&'a FrozenModule, FrozenError> {
+    vm.state.frozen.get(name).ok_or(FrozenError::NotFound)
+}
+
 #[pymodule]
 mod _imp {
     use crate::{
-        builtins::{PyBytesRef, PyCode, PyModule, PyStrRef},
+        builtins::{PyBytesRef, PyCode, PyMemoryView, PyModule, PyStrRef},
+        function::OptionalArg,
         import, PyObjectRef, PyRef, PyResult, TryFromObject, VirtualMachine,
     };
 
@@ -126,6 +156,30 @@ mod _imp {
         // TODO:
     }
 
+    #[allow(clippy::type_complexity)]
+    #[pyfunction]
+    fn find_frozen(
+        name: PyStrRef,
+        withdata: OptionalArg<bool>,
+        vm: &VirtualMachine,
+    ) -> PyResult<Option<(Option<PyRef<PyMemoryView>>, bool, PyStrRef)>> {
+        use super::FrozenError::*;
+
+        if withdata.into_option().is_some() {
+            // this is keyword-only argument in CPython
+            unimplemented!();
+        }
+
+        let info = match super::find_frozen(name.as_str(), vm) {
+            Ok(info) => info,
+            Err(NotFound | Disabled | BadName) => return Ok(None),
+            Err(e) => return Err(e.to_pyexception(name.as_str(), vm)),
+        };
+
+        let origname = name; // FIXME: origname != name
+        Ok(Some((None, info.package, origname)))
+    }
+
     #[pyfunction]
     fn source_hash(key: u64, source: PyBytesRef) -> Vec<u8> {
         let hash: u64 = crate::common::hash::keyed_hash(key, source.as_bytes());
