diff --git a/tests/snippets/dismod.py b/tests/snippets/dismod.py index cc9c1f63b8..eac1d3f726 100644 --- 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 index 5e1a30b033..d726b71106 100644 --- 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) }) }