Merge pull request #4031 from youknowone/encodings-feature

youknowone · web-flow · commit 662910e0c6ad · 2022-08-11T07:13:55.000+09:00
encodings as vm feature
diff --git a/Cargo.toml b/Cargo.toml
@@ -19,7 +19,9 @@ members = [
 ]
 
 [features]
-default = ["threading", "pylib", "stdlib", "zlib"]
+default = ["threading", "pylib", "stdlib", "zlib", "importlib", "encodings"]
+importlib = ["rustpython-vm/importlib"]
+encodings = ["rustpython-vm/encodings"]
 stdlib = ["rustpython-stdlib"]
 flame-it = ["rustpython-vm/flame-it", "flame", "flamescope"]
 freeze-stdlib = ["rustpython-vm/freeze-stdlib"]
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -10,6 +10,8 @@ include = ["src/**/*.rs", "Cargo.toml", "build.rs", "Lib/**/*.py"]
 
 [features]
 default = ["compile-parse", "pylib"]
+importlib = []
+encodings = ["importlib"]
 vm-tracing-logging = []
 flame-it = ["flame", "flamer"]
 pylib = ["rustpython-pylib"]
diff --git a/vm/Lib/core_modules/codecs.py b/vm/Lib/core_modules/codecs.py
@@ -0,0 +1 @@
+../../../Lib/codecs.py
diff --git a/vm/Lib/core_modules/encodings_utf_8.py b/vm/Lib/core_modules/encodings_utf_8.py
@@ -0,0 +1 @@
+../../../Lib/encodings/utf_8.py
diff --git a/vm/src/codecs.rs b/vm/src/codecs.rs
@@ -210,14 +210,25 @@ impl CodecsRegistry {
         Ok(())
     }
 
+    #[cfg(not(feature = "encodings"))]
+    pub(crate) fn register_manual(&self, name: &str, codec: PyCodec) -> PyResult<()> {
+        self.inner
+            .write()
+            .search_cache
+            .insert(name.to_owned(), codec);
+        Ok(())
+    }
+
     pub fn lookup(&self, encoding: &str, vm: &VirtualMachine) -> PyResult<PyCodec> {
         let encoding = normalize_encoding_name(encoding);
-        let inner = self.inner.read();
-        if let Some(codec) = inner.search_cache.get(encoding.as_ref()) {
-            return Ok(codec.clone());
-        }
-        let search_path = inner.search_path.clone();
-        drop(inner); // don't want to deadlock
+        let search_path = {
+            let inner = self.inner.read();
+            if let Some(codec) = inner.search_cache.get(encoding.as_ref()) {
+                // hit cache
+                return Ok(codec.clone());
+            }
+            inner.search_path.clone()
+        };
         let encoding = PyStr::from(encoding.into_owned()).into_ref(vm);
         for func in search_path {
             let res = vm.invoke(&func, (encoding.clone(),))?;
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -153,7 +153,8 @@ pub fn import_codeobj(
     sys_modules.set_item(module_name, module.clone(), vm)?;
 
     // Execute main code in module:
-    vm.run_code_obj(code_obj, Scope::with_builtins(None, attrs, vm))?;
+    let scope = Scope::with_builtins(None, attrs, vm);
+    vm.run_code_obj(code_obj, scope)?;
     Ok(module)
 }
 
diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs
@@ -187,6 +187,36 @@ impl VirtualMachine {
         vm
     }
 
+    /// set up the encodings search function
+    /// init_importlib must be called before this call
+    #[cfg(feature = "encodings")]
+    fn import_encodings(&mut self) -> PyResult<()> {
+        self.import("encodings", None, 0).map_err(|import_err| {
+            let err = self.new_runtime_error(
+                "Could not import encodings. Is your RUSTPYTHONPATH set? If you don't have \
+                    access to a consistent external environment (e.g. if you're embedding \
+                    rustpython in another application), try enabling the freeze-stdlib feature"
+                    .to_owned(),
+            );
+            err.set_cause(Some(import_err));
+            err
+        })?;
+        Ok(())
+    }
+
+    /// init only utf-8 encoding
+    #[cfg(not(feature = "encodings"))]
+    fn import_encodings(&mut self) -> PyResult<()> {
+        import::import_frozen(self, "codecs")?;
+        let encoding_module = import::import_frozen(self, "encodings_utf_8")?;
+        let getregentry = encoding_module.get_attr("getregentry", self)?;
+        let codec_info = self.invoke(&getregentry, ())?;
+        self.state
+            .codec_registry
+            .register_manual("utf-8", codec_info.try_into_value(self)?)?;
+        Ok(())
+    }
+
     fn initialize(&mut self) {
         flame_guard!("init VirtualMachine");
 
@@ -202,17 +232,7 @@ impl VirtualMachine {
             import::import_builtin(self, "_signal")?;
             import::init_importlib(self, self.state.settings.allow_external_library)?;
 
-            // set up the encodings search function
-            self.import("encodings", None, 0).map_err(|import_err| {
-                let err = self.new_runtime_error(
-                    "Could not import encodings. Is your RUSTPYTHONPATH set? If you don't have \
-                     access to a consistent external environment (e.g. if you're embedding \
-                     rustpython in another application), try enabling the freeze-stdlib feature"
-                        .to_owned(),
-                );
-                err.set_cause(Some(import_err));
-                err
-            })?;
+            self.import_encodings()?;
 
             #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))]
             {
diff --git a/vm/src/vm/setting.rs b/vm/src/vm/setting.rs
@@ -93,7 +93,7 @@ impl Default for Settings {
             argv: vec![],
             hash_seed: None,
             stdio_unbuffered: false,
-            allow_external_library: true,
+            allow_external_library: cfg!(feature = "importlib"),
         }
     }
 }
diff --git a/wasm/lib/Cargo.toml b/wasm/lib/Cargo.toml
@@ -19,7 +19,7 @@ no-start-func = []
 rustpython-parser = { path = "../../parser" }
 rustpython-common = { path = "../../common" }
 # make sure no threading! otherwise wasm build will fail
-rustpython-vm = { path = "../../vm", default-features = false, features = ["compile-parse"]  }
+rustpython-vm = { path = "../../vm", default-features = false, features = ["compile-parse", "encodings"] }
 rustpython-stdlib = { path = "../../stdlib", default-features = false, features = ["compile-parse"], optional = true }
 wasm-bindgen = "0.2"
 wasm-bindgen-futures = "0.4"

Original file line number	Diff line number	Diff line change
`@@ -153,7 +153,8 @@ pub fn import_codeobj(`
`153`	`153`	`sys_modules.set_item(module_name, module.clone(), vm)?;`
`154`	`154`
`155`	`155`	`// Execute main code in module:`
`156`		`- vm.run_code_obj(code_obj, Scope::with_builtins(None, attrs, vm))?;`
	`156`	`+ let scope = Scope::with_builtins(None, attrs, vm);`
	`157`	`+ vm.run_code_obj(code_obj, scope)?;`
`157`	`158`	`Ok(module)`
`158`	`159`	`}`
`159`	`160`
Original file line number	Diff line number	Diff line change
`@@ -93,7 +93,7 @@ impl Default for Settings {`
`93`	`93`	`argv: vec![],`
`94`	`94`	`hash_seed: None,`
`95`	`95`	`stdio_unbuffered: false,`
`96`		`- allow_external_library: true,`
	`96`	`+ allow_external_library: cfg!(feature = "importlib"),`
`97`	`97`	`}`
`98`	`98`	`}`
`99`	`99`	`}`