Merge pull request #4555 from youknowone/imp-override-frozen-modules

youknowone · web-flow · commit 6404c4bca8f7 · 2023-02-23T23:28:08.000+09:00
Add _imp._override_frozen_modules_for_tests
diff --git a/vm/src/stdlib/imp.rs b/vm/src/stdlib/imp.rs
@@ -116,6 +116,11 @@ mod _imp {
             .ok_or_else(|| vm.new_import_error(format!("No such frozen object named {name}"), name))
     }
 
+    #[pyfunction]
+    fn _override_frozen_modules_for_tests(value: isize, vm: &VirtualMachine) {
+        vm.state.override_frozen_modules.store(value);
+    }
+
     #[pyfunction]
     fn _fix_co_filename(_code: PyObjectRef, _path: PyStrRef) {
         // TODO:
diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs
@@ -96,6 +96,7 @@ pub struct PyGlobalState {
     pub codec_registry: CodecsRegistry,
     pub finalizing: AtomicBool,
     pub warnings: WarningsState,
+    pub override_frozen_modules: AtomicCell<isize>,
 }
 
 pub fn process_hash_secret_seed() -> u32 {
@@ -172,6 +173,7 @@ impl VirtualMachine {
                 codec_registry,
                 finalizing: AtomicBool::new(false),
                 warnings,
+                override_frozen_modules: AtomicCell::new(0),
             }),
             initialized: false,
             recursion_depth: Cell::new(0),
