Issue 9147: Add dis.code_info()

2 files changed, 200 insertions, 30 deletions

diff --git a/Lib/dis.py b/Lib/dis.py
index 2f46774..b4159e0 100644
--- a/Lib/dis.py
+++ b/Lib/dis.py
@@ -19,9 +19,6 @@ def _try_compile(source, name):
        Utility function to accept strings in functions that otherwise
        expect code objects
     """
-    # ncoghlan: currently only used by dis(), but plan to add an
-    # equivalent for show_code() as well (but one that returns a
-    # string rather than printing directly to the console)
     try:
         c = compile(source, name, 'eval')
     except SyntaxError:
@@ -37,11 +34,11 @@ def dis(x=None):
     if x is None:
         distb()
         return
-    if hasattr(x, '__func__'):
+    if hasattr(x, '__func__'):  # Method
         x = x.__func__
-    if hasattr(x, '__code__'):
+    if hasattr(x, '__code__'):  # Function
         x = x.__code__
-    if hasattr(x, '__dict__'):
+    if hasattr(x, '__dict__'):  # Class or module
         items = sorted(x.__dict__.items())
         for name, x1 in items:
             if isinstance(x1, _have_code):
@@ -51,11 +48,11 @@ def dis(x=None):
                 except TypeError as msg:
                     print("Sorry:", msg)
                 print()
-    elif hasattr(x, 'co_code'):
+    elif hasattr(x, 'co_code'): # Code object
         disassemble(x)
-    elif isinstance(x, (bytes, bytearray)):
+    elif isinstance(x, (bytes, bytearray)): # Raw bytecode
         _disassemble_bytes(x)
-    elif isinstance(x, str):
+    elif isinstance(x, str):    # Source code
         _disassemble_str(x)
     else:
         raise TypeError("don't know how to disassemble %s objects" %
@@ -97,35 +94,54 @@ def pretty_flags(flags):
         names.append(hex(flags))
     return ", ".join(names)
 
-def show_code(co):
-    """Show details about a code object."""
-    print("Name:             ", co.co_name)
-    print("Filename:         ", co.co_filename)
-    print("Argument count:   ", co.co_argcount)
-    print("Kw-only arguments:", co.co_kwonlyargcount)
-    print("Number of locals: ", co.co_nlocals)
-    print("Stack size:       ", co.co_stacksize)
-    print("Flags:            ", pretty_flags(co.co_flags))
+def code_info(x):
+    """Formatted details of methods, functions, or code."""
+    if hasattr(x, '__func__'): # Method
+        x = x.__func__
+    if hasattr(x, '__code__'): # Function
+        x = x.__code__
+    if isinstance(x, str):     # Source code
+        x = _try_compile(x, "<code_info>")
+    if hasattr(x, 'co_code'):  # Code object
+        return _format_code_info(x)
+    else:
+        raise TypeError("don't know how to disassemble %s objects" %
+                        type(x).__name__)
+
+def _format_code_info(co):
+    lines = []
+    lines.append("Name:              %s" % co.co_name)
+    lines.append("Filename:          %s" % co.co_filename)
+    lines.append("Argument count:    %s" % co.co_argcount)
+    lines.append("Kw-only arguments: %s" % co.co_kwonlyargcount)
+    lines.append("Number of locals:  %s" % co.co_nlocals)
+    lines.append("Stack size:        %s" % co.co_stacksize)
+    lines.append("Flags:             %s" % pretty_flags(co.co_flags))
     if co.co_consts:
-        print("Constants:")
+        lines.append("Constants:")
         for i_c in enumerate(co.co_consts):
-            print("%4d: %r" % i_c)
+            lines.append("%4d: %r" % i_c)
     if co.co_names:
-        print("Names:")
+        lines.append("Names:")
         for i_n in enumerate(co.co_names):
-            print("%4d: %s" % i_n)
+            lines.append("%4d: %s" % i_n)
     if co.co_varnames:
-        print("Variable names:")
+        lines.append("Variable names:")
         for i_n in enumerate(co.co_varnames):
-            print("%4d: %s" % i_n)
+            lines.append("%4d: %s" % i_n)
     if co.co_freevars:
-        print("Free variables:")
+        lines.append("Free variables:")
         for i_n in enumerate(co.co_freevars):
-            print("%4d: %s" % i_n)
+            lines.append("%4d: %s" % i_n)
     if co.co_cellvars:
-        print("Cell variables:")
+        lines.append("Cell variables:")
         for i_n in enumerate(co.co_cellvars):
-            print("%4d: %s" % i_n)
+            lines.append("%4d: %s" % i_n)
+    return "\n".join(lines)
+
+def show_code(co):
+    """Show details about a code object."""
+    print(code_info(co))
 
 def disassemble(co, lasti=-1):
     """Disassemble a code object."""
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index 7feee64..8f1783d 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -1,6 +1,6 @@
 # Minimal tests for dis module
 
-from test.support import run_unittest
+from test.support import run_unittest, captured_stdout
 import unittest
 import sys
 import dis
@@ -211,8 +211,162 @@ class DisTests(unittest.TestCase):
         self.do_disassembly_test(simple_stmt_str, dis_simple_stmt_str)
         self.do_disassembly_test(compound_stmt_str, dis_compound_stmt_str)
 
+code_info_code_info = """\
+Name:              code_info
+Filename:          {0}
+Argument count:    1
+Kw-only arguments: 0
+Number of locals:  1
+Stack size:        4
+Flags:             OPTIMIZED, NEWLOCALS, NOFREE
+Constants:
+   0: 'Formatted details of methods, functions, or code.'
+   1: '__func__'
+   2: '__code__'
+   3: '<code_info>'
+   4: 'co_code'
+   5: "don't know how to disassemble %s objects"
+   6: None
+Names:
+   0: hasattr
+   1: __func__
+   2: __code__
+   3: isinstance
+   4: str
+   5: _try_compile
+   6: _format_code_info
+   7: TypeError
+   8: type
+   9: __name__
+Variable names:
+   0: x""".format(dis.__file__)
+
+@staticmethod
+def tricky(x, y, z=True, *args, c, d, e=[], **kwds):
+    def f(c=c):
+        print(x, y, z, c, d, e, f)
+    yield x, y, z, c, d, e, f
+
+co_tricky_nested_f = tricky.__func__.__code__.co_consts[1]
+
+code_info_tricky = """\
+Name:              tricky
+Filename:          {0}
+Argument count:    3
+Kw-only arguments: 3
+Number of locals:  8
+Stack size:        7
+Flags:             OPTIMIZED, NEWLOCALS, VARARGS, VARKEYWORDS, GENERATOR
+Constants:
+   0: None
+   1: <code object f at {1}, file "{0}", line {2}>
+Variable names:
+   0: x
+   1: y
+   2: z
+   3: c
+   4: d
+   5: e
+   6: args
+   7: kwds
+Cell variables:
+   0: e
+   1: d
+   2: f
+   3: y
+   4: x
+   5: z""".format(__file__,
+                  hex(id(co_tricky_nested_f)),
+                  co_tricky_nested_f.co_firstlineno)
+
+code_info_tricky_nested_f = """\
+Name:              f
+Filename:          {0}
+Argument count:    1
+Kw-only arguments: 0
+Number of locals:  1
+Stack size:        8
+Flags:             OPTIMIZED, NEWLOCALS, NESTED
+Constants:
+   0: None
+Names:
+   0: print
+Variable names:
+   0: c
+Free variables:
+   0: e
+   1: d
+   2: f
+   3: y
+   4: x
+   5: z""".format(__file__)
+
+code_info_expr_str = """\
+Name:              <module>
+Filename:          <code_info>
+Argument count:    0
+Kw-only arguments: 0
+Number of locals:  0
+Stack size:        2
+Flags:             NOFREE
+Constants:
+   0: 1
+Names:
+   0: x"""
+
+code_info_simple_stmt_str = """\
+Name:              <module>
+Filename:          <code_info>
+Argument count:    0
+Kw-only arguments: 0
+Number of locals:  0
+Stack size:        2
+Flags:             NOFREE
+Constants:
+   0: 1
+   1: None
+Names:
+   0: x"""
+
+code_info_compound_stmt_str = """\
+Name:              <module>
+Filename:          <code_info>
+Argument count:    0
+Kw-only arguments: 0
+Number of locals:  0
+Stack size:        2
+Flags:             NOFREE
+Constants:
+   0: 0
+   1: 1
+   2: None
+Names:
+   0: x"""
+
+class CodeInfoTests(unittest.TestCase):
+    test_pairs = [
+      (dis.code_info, code_info_code_info),
+      (tricky, code_info_tricky),
+      (co_tricky_nested_f, code_info_tricky_nested_f),
+      (expr_str, code_info_expr_str),
+      (simple_stmt_str, code_info_simple_stmt_str),
+      (compound_stmt_str, code_info_compound_stmt_str),
+    ]
+
+    def test_code_info(self):
+        self.maxDiff = 1000
+        for x, expected in self.test_pairs:
+            self.assertEqual(dis.code_info(x), expected)
+
+    def test_show_code(self):
+        self.maxDiff = 1000
+        for x, expected in self.test_pairs:
+            with captured_stdout() as output:
+                dis.show_code(x)
+            self.assertEqual(output.getvalue(), expected+"\n")
+
 def test_main():
-    run_unittest(DisTests)
+    run_unittest(DisTests, CodeInfoTests)
 
 if __name__ == "__main__":
     test_main()