Add support for pulling __code__ out of a function (and a test that hits lots of instructions).

cthulahoops · cthulahoops · commit 79f43ad76252 · 2019-02-26T09:41:39.000Z
diff --git a/tests/snippets/dismod.py b/tests/snippets/dismod.py
@@ -8,3 +8,13 @@
 print("\n")
 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)
diff --git a/vm/src/stdlib/dis.rs b/vm/src/stdlib/dis.rs
@@ -3,7 +3,14 @@ use crate::pyobject::{PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol
 use crate::vm::VirtualMachine;
 
 fn dis_dis(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    dis_disassemble(vm, args)
+    arg_check!(vm, args, required = [(obj, None)]);
+    let code_name = vm.new_str("__code__".to_string());
+    let code = match vm.get_attribute(obj.clone(), code_name) {
+        Ok(co) => co,
+        Err(..) => obj.clone(),
+    };
+
+    dis_disassemble(vm, PyFuncArgs::new(vec![code], vec![]))
 }
 
 fn dis_disassemble(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
