Fix mini_repl and make ctx.new_code_object() more ergonomic

coolreader18 · coolreader18 · commit e74a8ec9fbec · 2020-10-19T23:55:55.000-05:00
diff --git a/examples/mini_repl.rs b/examples/mini_repl.rs
@@ -15,36 +15,17 @@ use vm::pyobject::ItemProtocol;
 // the program you're embedding this into doesn't take longer to start up.
 macro_rules! add_python_function {
     ( $scope:ident, $vm:ident, $src:literal $(,)? ) => {{
-        // this has to be in scope to turn a PyValue into a PyRef
-        // (a PyRef is a special reference that points to something in the VirtualMachine)
-        use vm::pyobject::PyValue;
-
+        // compile the code to bytecode
         let code = vm::py_compile!(source = $src);
+        // convert the rustpython_bytecode::CodeObject to a PyCodeRef
+        let code = $vm.ctx.new_code_object(code);
 
-        // takes the first constant in the file that's a function
-        let def = code
-            .get_constants()
-            .find_map(|c| match c {
-                vm::bytecode::Constant::Code { code } => Some(code),
-                _ => None,
-            })
-            .expect("No functions found in the provided module!");
-
-        // inserts the first function found in the module into the provided scope.
-        $scope.globals.set_item(
-            def.obj_name.as_str(),
-            $vm.ctx.new_pyfunction(
-                vm::builtins::PyCode::new(*def.clone()).into_ref(&$vm),
-                $scope.clone(),
-                None,
-                None,
-            ),
-            &$vm,
-        )
+        // run the python code in the scope to store the function
+        $vm.run_code_obj(code, $scope.clone())
     }};
 }
 
-static ON: AtomicBool = AtomicBool::new(false);
+static ON: AtomicBool = AtomicBool::new(true);
 
 fn on(b: bool) {
     ON.store(b, Ordering::Relaxed);
@@ -69,7 +50,10 @@ fn run(vm: &vm::VirtualMachine) -> vm::pyobject::PyResult<()> {
         vm,
         // a fun line to test this with is
         // ''.join( l * fib(i) for i, l in enumerate('supercalifragilistic') )
-        r#"def fib(n): return n if n <= 1 else fib(n - 1) + fib(n - 2)"#
+        r#"\
+def fib(n):
+    return n if n <= 1 else fib(n - 1) + fib(n - 2)
+"#
     )?;
 
     while ON.load(Ordering::Relaxed) {
diff --git a/vm/src/builtins/code.rs b/vm/src/builtins/code.rs
@@ -90,6 +90,20 @@ pub type PyCodeRef = PyRef<PyCode>;
 pub type CodeObject = bytecode::CodeObject<PyConstant>;
 pub type FrozenModule = bytecode::FrozenModule<PyConstant>;
 
+pub trait IntoCodeObject {
+    fn into_codeobj(self, ctx: &PyContext) -> CodeObject;
+}
+impl IntoCodeObject for CodeObject {
+    fn into_codeobj(self, _ctx: &PyContext) -> CodeObject {
+        self
+    }
+}
+impl IntoCodeObject for bytecode::CodeObject {
+    fn into_codeobj(self, ctx: &PyContext) -> CodeObject {
+        ctx.map_codeobj(self)
+    }
+}
+
 #[pyclass(module = false, name = "code")]
 pub struct PyCode {
     pub code: CodeObject,
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -346,8 +346,12 @@ impl PyContext {
         )
     }
 
-    pub fn new_code_object(&self, code: code::CodeObject) -> PyCodeRef {
-        PyRef::new_ref(code::PyCode::new(code), self.types.code_type.clone(), None)
+    pub fn new_code_object(&self, code: impl code::IntoCodeObject) -> PyCodeRef {
+        PyRef::new_ref(
+            code::PyCode::new(code.into_codeobj(self)),
+            self.types.code_type.clone(),
+            None,
+        )
     }
 
     pub fn new_pyfunction(
@@ -389,9 +393,7 @@ impl PyContext {
             bytecode::ConstantData::Str { value } => self.new_str(value),
             bytecode::ConstantData::Bytes { value } => self.new_bytes(value),
             bytecode::ConstantData::Boolean { value } => self.new_bool(value),
-            bytecode::ConstantData::Code { code } => {
-                self.new_code_object(self.map_codeobj(*code)).into_object()
-            }
+            bytecode::ConstantData::Code { code } => self.new_code_object(*code).into_object(),
             bytecode::ConstantData::Tuple { elements } => {
                 let elements = elements
                     .into_iter()
