From f752c93c0aa3a65e7b1b097e168cdb9bee44cf76 Mon Sep 17 00:00:00 2001 From: Windel Bouwman Date: Sun, 14 Jul 2019 12:33:22 +0200 Subject: [PATCH] Add sys.settrace and sys.setprofile functions to register profiling functions. --- tests/snippets/sysmod.py | 17 ++++++++++ vm/src/sysmodule.rs | 39 +++++++++++++++++++++++ vm/src/vm.rs | 67 ++++++++++++++++++++++++++++++++++++++-- 3 files changed, 121 insertions(+), 2 deletions(-) diff --git a/tests/snippets/sysmod.py b/tests/snippets/sysmod.py index 7c97846f1..0f81efce1 100644 --- a/tests/snippets/sysmod.py +++ b/tests/snippets/sysmod.py @@ -1,5 +1,6 @@ import sys +print('python executable:', sys.executable) print(sys.argv) assert sys.argv[0].endswith('.py') @@ -21,3 +22,19 @@ assert type(sys.flags).__name__ == "flags" assert type(sys.flags.optimize) is int assert sys.flags[3] == sys.flags.optimize assert sys.maxunicode == 1114111 + + +# Tracing: + +def trc(frame, event, arg): + print('trace event:', frame, event, arg) + +def demo(x): + print(x) + if x > 0: + demo(x - 1) + +sys.settrace(trc) +demo(5) +sys.settrace(None) + diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs index 10aeb4747..5f0c71996 100644 --- a/vm/src/sysmodule.rs +++ b/vm/src/sysmodule.rs @@ -18,6 +18,14 @@ fn argv(ctx: &PyContext) -> PyObjectRef { ctx.new_list(argv) } +fn executable(ctx: &PyContext) -> PyObjectRef { + if let Some(arg) = env::args().next() { + ctx.new_str(arg) + } else { + ctx.none() + } +} + fn getframe(offset: OptionalArg, vm: &VirtualMachine) -> PyResult { let offset = offset.into_option().unwrap_or(0); if offset > vm.frames.borrow().len() - 1 { @@ -110,6 +118,32 @@ fn sys_getfilesystemencodeerrors(_vm: &VirtualMachine) -> String { "surrogatepass".to_string() } +fn sys_getprofile(vm: &VirtualMachine) -> PyObjectRef { + vm.profile_func.borrow().clone() +} + +fn sys_setprofile(profilefunc: PyObjectRef, vm: &VirtualMachine) { + vm.profile_func.replace(profilefunc); + update_use_tracing(vm); +} + +fn sys_gettrace(vm: &VirtualMachine) -> PyObjectRef { + vm.trace_func.borrow().clone() +} + +fn sys_settrace(tracefunc: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef { + vm.trace_func.replace(tracefunc); + update_use_tracing(vm); + vm.ctx.none() +} + +fn update_use_tracing(vm: &VirtualMachine) { + let trace_is_none = vm.is_none(&vm.trace_func.borrow()); + let profile_is_none = vm.is_none(&vm.profile_func.borrow()); + let tracing = !(trace_is_none && profile_is_none); + vm.use_tracing.replace(tracing); +} + // TODO implement string interning, this will be key for performance fn sys_intern(value: PyStringRef, _vm: &VirtualMachine) -> PyStringRef { value @@ -264,12 +298,15 @@ settrace() -- set the global debug tracing function "builtin_module_names" => builtin_module_names, "byteorder" => ctx.new_str(bytorder), "copyright" => ctx.new_str(copyright.to_string()), + "executable" => executable(ctx), "flags" => flags, "getrefcount" => ctx.new_rustfunc(sys_getrefcount), "getsizeof" => ctx.new_rustfunc(sys_getsizeof), "implementation" => implementation, "getfilesystemencoding" => ctx.new_rustfunc(sys_getfilesystemencoding), "getfilesystemencodeerrors" => ctx.new_rustfunc(sys_getfilesystemencodeerrors), + "getprofile" => ctx.new_rustfunc(sys_getprofile), + "gettrace" => ctx.new_rustfunc(sys_gettrace), "intern" => ctx.new_rustfunc(sys_intern), "maxunicode" => ctx.new_int(0x0010_FFFF), "maxsize" => ctx.new_int(std::isize::MAX), @@ -286,6 +323,8 @@ settrace() -- set the global debug tracing function "path_importer_cache" => ctx.new_dict(), "pycache_prefix" => vm.get_none(), "dont_write_bytecode" => vm.new_bool(true), + "setprofile" => ctx.new_rustfunc(sys_setprofile), + "settrace" => ctx.new_rustfunc(sys_settrace), "version" => vm.new_str(version::get_version()), }); diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 232eca0e1..af528521f 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -7,6 +7,7 @@ use std::cell::{Ref, RefCell}; use std::collections::hash_map::HashMap; use std::collections::hash_set::HashSet; +use std::fmt; use std::rc::Rc; use std::sync::{Mutex, MutexGuard}; @@ -56,6 +57,9 @@ pub struct VirtualMachine { pub exceptions: RefCell>, pub frozen: RefCell>, pub import_func: RefCell, + pub profile_func: RefCell, + pub trace_func: RefCell, + pub use_tracing: RefCell, pub settings: PySettings, } @@ -86,6 +90,22 @@ pub struct PySettings { pub quiet: bool, } +/// Trace events for sys.settrace and sys.setprofile. +enum TraceEvent { + Call, + Return, +} + +impl fmt::Display for TraceEvent { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + use TraceEvent::*; + match self { + Call => write!(f, "call"), + Return => write!(f, "return"), + } + } +} + /// Sensible default settings. impl Default for PySettings { fn default() -> Self { @@ -115,6 +135,8 @@ impl VirtualMachine { let stdlib_inits = RefCell::new(stdlib::get_module_inits()); let frozen = RefCell::new(frozen::get_module_inits()); let import_func = RefCell::new(ctx.none()); + let profile_func = RefCell::new(ctx.none()); + let trace_func = RefCell::new(ctx.none()); let vm = VirtualMachine { builtins: builtins.clone(), sys_module: sysmod.clone(), @@ -125,6 +147,9 @@ impl VirtualMachine { exceptions: RefCell::new(vec![]), frozen, import_func, + profile_func, + trace_func, + use_tracing: RefCell::new(false), settings, }; @@ -323,6 +348,11 @@ impl VirtualMachine { self.ctx.none() } + /// Test whether a python object is `None`. + pub fn is_none(&self, obj: &PyObjectRef) -> bool { + obj.is(&self.get_none()) + } + pub fn get_type(&self) -> PyClassRef { self.ctx.type_type() } @@ -452,7 +482,10 @@ impl VirtualMachine { ref kw_only_defaults, }) = func_ref.payload() { - self.invoke_python_function(code, scope, defaults, kw_only_defaults, args) + self.trace_event(TraceEvent::Call)?; + let res = self.invoke_python_function(code, scope, defaults, kw_only_defaults, args); + self.trace_event(TraceEvent::Return)?; + res } else if let Some(PyMethod { ref function, ref object, @@ -474,7 +507,37 @@ impl VirtualMachine { where T: Into, { - self._invoke(func_ref, args.into()) + let res = self._invoke(func_ref, args.into()); + res + } + + /// Call registered trace function. + fn trace_event(&self, event: TraceEvent) -> PyResult<()> { + if *self.use_tracing.borrow() { + let frame = self.get_none(); + let event = self.new_str(event.to_string()); + let arg = self.get_none(); + let args = vec![frame, event, arg]; + + // temporarily disable tracing, during the call to the + // tracing function itself. + let trace_func = self.trace_func.borrow().clone(); + if !self.is_none(&trace_func) { + self.use_tracing.replace(false); + let res = self.invoke(trace_func, args.clone()); + self.use_tracing.replace(true); + res?; + } + + let profile_func = self.profile_func.borrow().clone(); + if !self.is_none(&profile_func) { + self.use_tracing.replace(false); + let res = self.invoke(profile_func, args); + self.use_tracing.replace(true); + res?; + } + } + Ok(()) } fn invoke_python_function(