dis.rs - small refactor dis_dis.

cthulahoops · cthulahoops · commit 73a65a1c779a · 2019-02-26T09:41:39.000Z
diff --git a/tests/snippets/dismod.py b/tests/snippets/dismod.py
@@ -18,3 +18,17 @@ def f():
             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
@@ -4,13 +4,14 @@ 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());
-    let code = match vm.get_attribute(obj.clone(), code_name) {
-        Ok(co) => co,
-        Err(..) => obj.clone(),
-    };
+    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![code], vec![]))
+    dis_disassemble(vm, PyFuncArgs::new(vec![obj.clone()], vec![]))
 }
 
 fn dis_disassemble(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
