Add extremely minimal dis module.

cthulahoops · cthulahoops · commit 5ebfd55b20d5 · 2019-02-21T14:02:59.000Z
diff --git a/tests/snippets/dismod.py b/tests/snippets/dismod.py
@@ -0,0 +1,10 @@
+import dis
+
+dis.disassemble(compile("5 + x + 5 or 2", "", "eval"))
+print("\n")
+dis.disassemble(compile("def f(x):\n   return 1", "", "exec"))
+print("\n")
+dis.disassemble(compile("if a:\n 1 or 2\nelif x == 'hello':\n 3\nelse:\n 4", "", "exec"))
+print("\n")
+dis.disassemble(compile("f(x=1, y=2)", "", "eval"))
+print("\n")
diff --git a/vm/src/stdlib/dis.rs b/vm/src/stdlib/dis.rs
@@ -0,0 +1,18 @@
+use super::super::obj::objcode;
+use super::super::obj::objtype;
+use super::super::pyobject::{PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol};
+use super::super::vm::VirtualMachine;
+
+fn dis_disassemble(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(co, Some(vm.ctx.code_type()))]);
+
+    let code = objcode::get_value(co);
+    print!("{}", code);
+    Ok(vm.get_none())
+}
+
+pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
+    let py_mod = ctx.new_module("dis", ctx.new_scope(None));
+    ctx.set_attr(&py_mod, "disassemble", ctx.new_rustfunc(dis_disassemble));
+    py_mod
+}
diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs
@@ -1,4 +1,5 @@
 mod ast;
+mod dis;
 pub mod io;
 mod json;
 mod keyword;
@@ -21,6 +22,7 @@ pub type StdlibInitFunc = fn(&PyContext) -> PyObjectRef;
 pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {
     let mut modules = HashMap::new();
     modules.insert("ast".to_string(), ast::mk_module as StdlibInitFunc);
+    modules.insert("dis".to_string(), dis::mk_module as StdlibInitFunc);
     modules.insert("io".to_string(), io::mk_module as StdlibInitFunc);
     modules.insert("json".to_string(), json::mk_module as StdlibInitFunc);
     modules.insert("keyword".to_string(), keyword::mk_module as StdlibInitFunc);
