RustPython · windelbouwman · Feb 21, 2019 · Feb 21, 2019 · Feb 21, 2019
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/bytecode.rs b/vm/src/bytecode.rs
@@ -8,7 +8,7 @@
 use num_bigint::BigInt;
 use num_complex::Complex64;
 use rustpython_parser::ast;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fmt;
 
 /// Primary container of a single code object. Each python function has
@@ -173,6 +173,8 @@ pub enum Instruction {
     },
 }
 
+use self::Instruction::*;
+
 #[derive(Debug, Clone, PartialEq)]
 pub enum CallType {
     Positional(usize),
@@ -277,17 +279,123 @@ impl CodeObject {
     }
 }
 
+impl fmt::Display for CodeObject {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let label_targets: HashSet<&usize> = self.label_map.values().collect();
+        for (offset, instruction) in self.instructions.iter().enumerate() {
+            let arrow = if label_targets.contains(&offset) {
+                ">>"
+            } else {
+                "  "
+            };
+            write!(f, "          {} {:5} ", arrow, offset)?;
+            instruction.fmt_dis(f, &self.label_map)?;
+        }
+        Ok(())
+    }
+}
+
+impl Instruction {
+    fn fmt_dis(&self, f: &mut fmt::Formatter, label_map: &HashMap<Label, usize>) -> fmt::Result {
+        macro_rules! w {
+            ($variant:ident) => {
+                write!(f, "{:20}\n", stringify!($variant))
+            };
+            ($variant:ident, $var:expr) => {
+                write!(f, "{:20} ({})\n", stringify!($variant), $var)
+            };
+            ($variant:ident, $var1:expr, $var2:expr) => {
+                write!(f, "{:20} ({}, {})\n", stringify!($variant), $var1, $var2)
+            };
+        }
+
+        match self {
+            Import { name, symbol } => w!(Import, name, format!("{:?}", symbol)),
+            ImportStar { name } => w!(ImportStar, name),
+            LoadName { name } => w!(LoadName, name),
+            StoreName { name } => w!(StoreName, name),
+            DeleteName { name } => w!(DeleteName, name),
+            StoreSubscript => w!(StoreSubscript),
+            DeleteSubscript => w!(DeleteSubscript),
+            StoreAttr { name } => w!(StoreAttr, name),
+            DeleteAttr { name } => w!(DeleteAttr, name),
+            LoadConst { value } => w!(LoadConst, value),
+            UnaryOperation { op } => w!(UnaryOperation, format!("{:?}", op)),
+            BinaryOperation { op, inplace } => w!(BinaryOperation, format!("{:?}", op), inplace),
+            LoadAttr { name } => w!(LoadAttr, name),
+            CompareOperation { op } => w!(CompareOperation, format!("{:?}", op)),
+            Pop => w!(Pop),
+            Rotate { amount } => w!(Rotate, amount),
+            Duplicate => w!(Duplicate),
+            GetIter => w!(GetIter),
+            Pass => w!(Pass),
+            Continue => w!(Continue),
+            Break => w!(Break),
+            Jump { target } => w!(Jump, label_map[target]),
+            JumpIf { target } => w!(JumpIf, label_map[target]),
+            JumpIfFalse { target } => w!(JumpIfFalse, label_map[target]),
+            MakeFunction { flags } => w!(MakeFunction, format!("{:?}", flags)),
+            CallFunction { typ } => w!(CallFunction, format!("{:?}", typ)),
+            ForIter { target } => w!(ForIter, label_map[target]),
+            ReturnValue => w!(ReturnValue),
+            YieldValue => w!(YieldValue),
+            YieldFrom => w!(YieldFrom),
+            SetupLoop { start, end } => w!(SetupLoop, label_map[start], label_map[end]),
+            SetupExcept { handler } => w!(SetupExcept, handler),
+            SetupWith { end } => w!(SetupWith, end),
+            CleanupWith { end } => w!(CleanupWith, end),
+            PopBlock => w!(PopBlock),
+            Raise { argc } => w!(Raise, argc),
+            BuildString { size } => w!(BuildString, size),
+            BuildTuple { size, unpack } => w!(BuildTuple, size, unpack),
+            BuildList { size, unpack } => w!(BuildList, size, unpack),
+            BuildSet { size, unpack } => w!(BuildSet, size, unpack),
+            BuildMap { size, unpack } => w!(BuildMap, size, unpack),
+            BuildSlice { size } => w!(BuildSlice, size),
+            ListAppend { i } => w!(ListAppend, i),
+            SetAdd { i } => w!(SetAdd, i),
+            MapAdd { i } => w!(MapAdd, i),
+            PrintExpr => w!(PrintExpr),
+            LoadBuildClass => w!(LoadBuildClass),
+            StoreLocals => w!(StoreLocals),
+            UnpackSequence { size } => w!(UnpackSequence, size),
+            UnpackEx { before, after } => w!(UnpackEx, before, after),
+            Unpack => w!(Unpack),
+            FormatValue { spec } => w!(FormatValue, spec),
+        }
+    }
+}
+
+impl fmt::Display for Constant {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Constant::Integer { value } => write!(f, "{}", value),
+            Constant::Float { value } => write!(f, "{}", value),
+            Constant::Complex { value } => write!(f, "{}", value),
+            Constant::Boolean { value } => write!(f, "{}", value),
+            Constant::String { value } => write!(f, "{:?}", value),
+            Constant::Bytes { value } => write!(f, "{:?}", value),
+            Constant::Code { code } => write!(f, "{:?}", code),
+            Constant::Tuple { elements } => write!(
+                f,
+                "({})",
+                elements
+                    .iter()
+                    .map(|e| format!("{}", e))
+                    .collect::<Vec<_>>()
+                    .join(", ")
+            ),
+            Constant::None => write!(f, "None"),
+        }
+    }
+}
+
 impl fmt::Debug for CodeObject {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let inst_str = self
-            .instructions
-            .iter()
-            .zip(self.locations.iter())
-            .enumerate()
-            .map(|(i, inst)| format!("Inst {}: {:?}", i, inst))
-            .collect::<Vec<_>>()
-            .join("\n");
-        let labelmap_str = format!("label_map: {:?}", self.label_map);
-        write!(f, "Code Object {{ \n{}\n{} }}", inst_str, labelmap_str)
+        write!(
+            f,
+            "<code object {} at ??? file {:?}, line {}>",
+            self.obj_name, self.source_path, self.first_line_number
+        )
     }
 }
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);