Merge pull request #2054 from BolunThompson/str_dis

youknowone · web-flow · commit 55e97af2f4de · 2020-08-02T00:40:57.000+09:00
Change dis.dis to take strings to decompile
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -0,0 +1,60 @@
+import subprocess
+import sys
+import unittest
+
+# This only tests that it prints something in order
+# to avoid changing this test if the bytecode changes
+
+# These tests start a new process instead of redirecting stdout because
+# stdout is being written to by rust code, which currently can't be
+# redirected by reassigning sys.stdout
+
+
+class TestDis(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.setup = """
+import dis
+def tested_func(): pass
+"""
+        cls.command = (sys.executable, "-c")
+
+    def test_dis(self):
+        test_code = f"""
+{self.setup}
+dis.dis(tested_func)
+dis.dis("x = 2; print(x)")
+"""
+
+        result = subprocess.run(
+            self.command + (test_code,), capture_output=True
+        )
+        self.assertNotEqual("", result.stdout.decode())
+        self.assertEqual("", result.stderr.decode())
+
+    def test_disassemble(self):
+        test_code = f"""
+{self.setup}
+dis.disassemble(tested_func)
+"""
+        result = subprocess.run(
+            self.command + (test_code,), capture_output=True
+        )
+        # In CPython this would raise an AttributeError, not a
+        # TypeError because dis is implemented in python in CPython and
+        # as such the type mismatch wouldn't be caught immeadiately
+        self.assertIn("TypeError", result.stderr.decode())
+
+        test_code = f"""
+{self.setup}
+dis.disassemble(tested_func.__code__)
+"""
+        result = subprocess.run(
+            self.command + (test_code,), capture_output=True
+        )
+        self.assertNotEqual("", result.stdout.decode())
+        self.assertEqual("", result.stderr.decode())
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/vm/src/stdlib/dis.rs b/vm/src/stdlib/dis.rs
@@ -1,14 +1,25 @@
 use crate::bytecode::CodeFlags;
 use crate::obj::objcode::PyCodeRef;
+use crate::obj::objstr::PyStringRef;
 use crate::pyobject::{ItemProtocol, PyObjectRef, PyResult, TryFromObject};
 use crate::vm::VirtualMachine;
+use rustpython_compiler::compile;
 
 fn dis_dis(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult {
     // Method or function:
     if let Ok(co) = vm.get_attribute(obj.clone(), "__code__") {
         return dis_disassemble(co, vm);
     }
 
+    // String:
+    if let Ok(co_str) = PyStringRef::try_from_object(vm, obj.clone()) {
+        let code = vm
+            .compile(co_str.as_str(), compile::Mode::Exec, "<string>".to_owned())
+            .map_err(|err| vm.new_syntax_error(&err))?
+            .into_object();
+        return dis_disassemble(code, vm);
+    }
+
     dis_disassemble(obj, vm)
 }
 
