Add sys.settrace and sys.setprofile functions to register profiling functions.

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)
+

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<usize>, vm: &VirtualMachine) -> PyResult<FrameRef> {
     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()),
     });
 

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<Vec<PyObjectRef>>,
     pub frozen: RefCell<HashMap<String, bytecode::CodeObject>>,
     pub import_func: RefCell<PyObjectRef>,
+    pub profile_func: RefCell<PyObjectRef>,
+    pub trace_func: RefCell<PyObjectRef>,
+    pub use_tracing: RefCell<bool>,
     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<PyFuncArgs>,
     {
-        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(