diff --git a/src/main.rs b/src/main.rs index db94cd15e..61c44cb50 100644 --- a/src/main.rs +++ b/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()) { diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs index 0eb46dc03..66258590d 100644 --- a/vm/src/sysmodule.rs +++ b/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 }