Implement sys.getrefcount, sys.ps1, sys.ps2 and sys.getsizeof

Author: windelbouwman

src/main.rs

@@ -12,7 +12,7 @@ use clap::{App, Arg};
 use rustpython_parser::parser;
 use rustpython_vm::obj::objstr;
 use rustpython_vm::print_exception;
-use rustpython_vm::pyobject::{PyObjectRef, PyResult};
+use rustpython_vm::pyobject::{AttributeProtocol, PyObjectRef, PyResult};
 use rustpython_vm::VirtualMachine;
 use rustpython_vm::{compile, import};
 use rustyline::error::ReadlineError;
@@ -167,7 +167,11 @@ fn run_shell(vm: &mut VirtualMachine) -> PyResult {
         //        Err(_) => ">>>>> ".to_string(),
         //};
 
-        match rl.readline(">>>>> ") {
+        // We can customize the prompt:
+        let ps1 = objstr::get_value(&vm.sys_module.get_attr("ps1").unwrap());
+        let ps2 = objstr::get_value(&vm.sys_module.get_attr("ps2").unwrap());
+
+        match rl.readline(&ps1) {
             Ok(line) => {
                 input.push_str(&line);
                 input.push_str("\n");
@@ -184,7 +188,7 @@ fn run_shell(vm: &mut VirtualMachine) -> PyResult {
                         //        Ok(value) => objstr::get_value(&value),
                         //        Err(_) => "..... ".to_string(),
                         //};
-                        match rl.readline("..... ") {
+                        match rl.readline(&ps2) {
                             Ok(line) => {
                                 if line.len() == 0 {
                                     if shell_exec(vm, &input, vars.clone()) {

vm/src/sysmodule.rs

@@ -1,6 +1,9 @@
-use super::pyobject::{DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult};
-use super::vm::VirtualMachine;
-use std::env;
+use num_bigint::ToBigInt;
+use obj::objtype;
+use pyobject::{DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol};
+use std::rc::Rc;
+use std::{env, mem};
+use vm::VirtualMachine;
 
 /*
  * The magic sys module.
@@ -20,6 +23,19 @@ fn getframe(vm: &mut VirtualMachine, _args: PyFuncArgs) -> PyResult {
     }
 }
 
+fn sys_getrefcount(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(object, None)]);
+    let size = Rc::strong_count(&object);
+    Ok(vm.ctx.new_int(size.to_bigint().unwrap()))
+}
+
+fn sys_getsizeof(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(object, None)]);
+    // TODO: implement default optional argument.
+    let size = mem::size_of_val(&object.borrow());
+    Ok(vm.ctx.new_int(size.to_bigint().unwrap()))
+}
+
 pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
     let path_list = match env::var_os("PYTHONPATH") {
         Some(paths) => env::split_paths(&paths)
@@ -34,7 +50,13 @@ pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
     modules.set_item(&sys_name, sys_mod.clone());
     sys_mod.set_item("modules", modules);
     sys_mod.set_item("argv", argv(ctx));
+    sys_mod.set_item("getrefcount", ctx.new_rustfunc(sys_getrefcount));
+    sys_mod.set_item("getsizeof", ctx.new_rustfunc(sys_getsizeof));
+    let maxsize = ctx.new_int(std::usize::MAX.to_bigint().unwrap());
+    sys_mod.set_item("maxsize", maxsize);
     sys_mod.set_item("path", path);
+    sys_mod.set_item("ps1", ctx.new_str(">>>>> ".to_string()));
+    sys_mod.set_item("ps2", ctx.new_str("..... ".to_string()));
     sys_mod.set_item("_getframe", ctx.new_rustfunc(getframe));
     sys_mod
 }