Merge pull request #555 from RustPython/dis_dis

windelbouwman · web-flow · commit c3309ee867eb · 2019-02-26T17:07:41.000+01:00
Dis dis
diff --git a/tests/snippets/dismod.py b/tests/snippets/dismod.py
@@ -1,10 +1,34 @@
 import dis
 
-dis.disassemble(compile("5 + x + 5 or 2", "", "eval"))
+dis.dis(compile("5 + x + 5 or 2", "", "eval"))
 print("\n")
-dis.disassemble(compile("def f(x):\n   return 1", "", "exec"))
+dis.dis(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"))
+dis.dis(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"))
+dis.dis(compile("f(x=1, y=2)", "", "eval"))
 print("\n")
+
+def f():
+    with g():
+        try:
+            for a in {1: 4, 2: 5}:
+                yield [True and False or True, []]
+        except Exception:
+            raise not ValueError({1 for i in [1,2,3]})
+
+dis.dis(f)
+
+class A(object):
+    def f():
+        x += 1
+        pass
+    def g():
+        for i in range(5):
+            if i:
+                continue
+            else:
+                break
+
+print("A.f\n")
+dis.dis(A.f)
diff --git a/vm/src/stdlib/dis.rs b/vm/src/stdlib/dis.rs
@@ -2,6 +2,18 @@ use crate::obj::objcode;
 use crate::pyobject::{PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol};
 use crate::vm::VirtualMachine;
 
+fn dis_dis(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, None)]);
+
+    // Method or function:
+    let code_name = vm.new_str("__code__".to_string());
+    if let Ok(co) = vm.get_attribute(obj.clone(), code_name) {
+        return dis_disassemble(vm, PyFuncArgs::new(vec![co], vec![]));
+    }
+
+    dis_disassemble(vm, PyFuncArgs::new(vec![obj.clone()], vec![]))
+}
+
 fn dis_disassemble(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(co, Some(vm.ctx.code_type()))]);
 
@@ -12,6 +24,7 @@ fn dis_disassemble(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 
 pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
     py_module!(ctx, "dis", {
+        "dis" => ctx.new_rustfunc(dis_dis),
         "disassemble" => ctx.new_rustfunc(dis_disassemble)
     })
 }
