From 0acce07631ea8fcdea4c7c100105eaeca4047690 Mon Sep 17 00:00:00 2001
From: Aviv Palivoda <palaviv@gmail.com>
Date: Sat, 1 Jun 2019 15:11:55 +0300
Subject: [PATCH] Add _imp.init_frozen

---
 tests/snippets/imp.py |  3 +++
 vm/src/stdlib/imp.rs  | 10 ++++++++++
 2 files changed, 13 insertions(+)

diff --git a/tests/snippets/imp.py b/tests/snippets/imp.py
index 0f60c8082..44e7e1042 100644
--- a/tests/snippets/imp.py
+++ b/tests/snippets/imp.py
@@ -23,3 +23,6 @@ assert _imp.create_builtin(B) == None
 _imp.exec_builtin(imp_time) == 0
 
 _imp.get_frozen_object("__hello__")
+
+hello = _imp.init_frozen("__hello__")
+assert hello.initialized == True
diff --git a/vm/src/stdlib/imp.rs b/vm/src/stdlib/imp.rs
index 6aa233b8d..ca266aa27 100644
--- a/vm/src/stdlib/imp.rs
+++ b/vm/src/stdlib/imp.rs
@@ -1,4 +1,5 @@
 use crate::compile;
+use crate::import::import_file;
 use crate::obj::objcode::PyCodeRef;
 use crate::obj::objmodule::PyModuleRef;
 use crate::obj::objstr;
@@ -63,6 +64,14 @@ fn imp_get_frozen_object(name: PyStringRef, vm: &VirtualMachine) -> PyResult<PyC
     }
 }
 
+fn imp_init_frozen(name: PyStringRef, vm: &VirtualMachine) -> PyResult {
+    if let Some(frozen) = vm.frozen.borrow().get(name.as_str()) {
+        import_file(vm, name.as_str(), "frozen".to_string(), frozen.to_string())
+    } else {
+        Err(vm.new_import_error(format!("No such frozen object named {}", name.as_str())))
+    }
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
     let module = py_module!(vm, "_imp", {
@@ -75,6 +84,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "create_builtin" => ctx.new_rustfunc(imp_create_builtin),
         "exec_builtin" => ctx.new_rustfunc(imp_exec_builtin),
         "get_frozen_object" => ctx.new_rustfunc(imp_get_frozen_object),
+        "init_frozen" => ctx.new_rustfunc(imp_init_frozen),
     });
 
     module