Implement sys._getframe() to the point where we can use it to get locals.

cthulahoops · cthulahoops · commit a1b7c61fcf40 · 2018-11-11T17:54:30.000Z
diff --git a/tests/snippets/getframe.py b/tests/snippets/getframe.py
@@ -0,0 +1,20 @@
+import sys
+
+value = 189
+locals_dict = sys._getframe().f_locals
+assert locals_dict['value'] == 189
+foo = 'bar'
+assert locals_dict['foo'] == foo
+
+def test_function():
+    x = 17
+    assert sys._getframe().f_locals is not locals_dict
+    assert sys._getframe().f_locals['x'] == 17
+
+test_function()
+
+class TestClass():
+    def __init__(self):
+        assert sys._getframe().f_locals['self'] == self
+
+TestClass()
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -3,6 +3,7 @@ extern crate rustpython_parser;
 use self::rustpython_parser::ast;
 use std::collections::hash_map::HashMap;
 use std::fmt;
+use std::mem;
 use std::path::PathBuf;
 
 use super::builtins;
@@ -38,6 +39,7 @@ enum Block {
     },
 }
 
+#[derive(Clone)]
 pub struct Frame {
     pub code: bytecode::CodeObject,
     // We need 1 stack per frame
@@ -94,6 +96,8 @@ impl Frame {
             "<unknown>".to_string()
         };
 
+        let prev_frame = mem::replace(&mut vm.current_frame, Some(vm.ctx.new_frame(self.clone())));
+
         // This is the name of the object being run:
         let run_obj_name = &self.code.obj_name.to_string();
 
@@ -142,6 +146,7 @@ impl Frame {
             }
         };
 
+        vm.current_frame = prev_frame;
         value
     }
 
diff --git a/vm/src/obj/objframe.rs b/vm/src/obj/objframe.rs
@@ -2,8 +2,9 @@
 
 */
 
+use super::super::frame::Frame;
 use super::super::pyobject::{
-    AttributeProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol,
+    AttributeProtocol, PyContext, PyFuncArgs, PyObjectKind, PyObjectRef, PyResult, TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
 use super::objtype;
@@ -12,6 +13,7 @@ pub fn init(context: &PyContext) {
     let ref frame_type = context.frame_type;
     frame_type.set_attr("__new__", context.new_rustfunc(frame_new));
     frame_type.set_attr("__repr__", context.new_rustfunc(frame_repr));
+    frame_type.set_attr("f_locals", context.new_property(frame_locals));
 }
 
 fn frame_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -24,3 +26,24 @@ fn frame_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let repr = format!("<frame object at .. >");
     Ok(vm.new_str(repr))
 }
+
+fn frame_locals(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(frame, Some(vm.ctx.frame_type()))]);
+    let frame = get_value(frame);
+    let py_scope = frame.locals.clone();
+    let py_scope = py_scope.borrow();
+
+    if let PyObjectKind::Scope { scope } = &py_scope.kind {
+        Ok(scope.locals.clone())
+    } else {
+        panic!("The scope isn't a scope!");
+    }
+}
+
+pub fn get_value(obj: &PyObjectRef) -> Frame {
+    if let PyObjectKind::Frame { frame } = &obj.borrow().kind {
+        frame.clone()
+    } else {
+        panic!("Inner error getting int {:?}", obj);
+    }
+}
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -459,6 +459,22 @@ impl PyContext {
         )
     }
 
+    pub fn new_frame(&self, frame: Frame) -> PyObjectRef {
+        PyObject::new(PyObjectKind::Frame { frame: frame }, self.frame_type())
+    }
+
+    pub fn new_property(&self, function: RustPyFunc) -> PyObjectRef {
+        let fget = self.new_rustfunc(function);
+        let py_obj = PyObject::new(
+            PyObjectKind::Instance {
+                dict: self.new_dict(),
+            },
+            self.property_type(),
+        );
+        py_obj.set_attr("fget", fget.clone());
+        py_obj
+    }
+
     pub fn new_function(
         &self,
         code_obj: PyObjectRef,
@@ -782,6 +798,9 @@ pub enum PyObjectKind {
     Code {
         code: bytecode::CodeObject,
     },
+    Frame {
+        frame: Frame,
+    },
     Function {
         code: PyObjectRef,
         scope: PyObjectRef,
@@ -856,6 +875,7 @@ impl fmt::Debug for PyObjectKind {
             } => write!(f, "class {:?}", name),
             &PyObjectKind::Instance { dict: _ } => write!(f, "instance"),
             &PyObjectKind::RustFunction { function: _ } => write!(f, "rust function"),
+            &PyObjectKind::Frame { .. } => write!(f, "frame"),
         }
     }
 }
@@ -913,6 +933,7 @@ impl PyObject {
             PyObjectKind::Code { code: _ } => format!("<code>"),
             PyObjectKind::Function { .. } => format!("<func>"),
             PyObjectKind::Generator { .. } => format!("<generator>"),
+            PyObjectKind::Frame { .. } => format!("<frame>"),
             PyObjectKind::BoundMethod { .. } => format!("<bound-method>"),
             PyObjectKind::RustFunction { function: _ } => format!("<rustfunc>"),
             PyObjectKind::Module { ref name, dict: _ } => format!("<module '{}'>", name),
diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs
@@ -1,4 +1,5 @@
-use super::pyobject::{DictProtocol, PyContext, PyObjectRef};
+use super::pyobject::{DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult};
+use super::vm::VirtualMachine;
 use std::env;
 
 /*
@@ -11,6 +12,14 @@ fn argv(ctx: &PyContext) -> PyObjectRef {
     ctx.new_list(argv)
 }
 
+fn getframe(vm: &mut VirtualMachine, _args: PyFuncArgs) -> PyResult {
+    if let Some(frame) = &vm.current_frame {
+        Ok(frame.clone())
+    } else {
+        panic!("Current frame is undefined!")
+    }
+}
+
 pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
     let path_list = match env::var_os("PYTHONPATH") {
         Some(paths) => env::split_paths(&paths)
@@ -23,8 +32,9 @@ pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
     let sys_name = "sys".to_string();
     let sys_mod = ctx.new_module(&sys_name, ctx.new_scope(None));
     modules.set_item(&sys_name, sys_mod.clone());
-    sys_mod.set_item(&"modules".to_string(), modules);
-    sys_mod.set_item(&"argv".to_string(), argv(ctx));
-    sys_mod.set_item(&"path".to_string(), path);
+    sys_mod.set_item("modules", modules);
+    sys_mod.set_item("argv", argv(ctx));
+    sys_mod.set_item("path", path);
+    sys_mod.set_item("_getframe", ctx.new_rustfunc(getframe));
     sys_mod
 }
