bpo-11822: Improve disassembly to show embedded code objects. (#1844)

The depth argument limits recursion.

5 files changed, 121 insertions, 15 deletions

diff --git a/Doc/library/dis.rst b/Doc/library/dis.rst
index f82dc40..bc32380 100644
--- a/Doc/library/dis.rst
+++ b/Doc/library/dis.rst
@@ -138,23 +138,32 @@ operation is being performed, so the intermediate analysis object isn't useful:
       Added *file* parameter.
 
 
-.. function:: dis(x=None, *, file=None)
+.. function:: dis(x=None, *, file=None, depth=None)
 
    Disassemble the *x* object.  *x* can denote either a module, a class, a
    method, a function, a generator, a code object, a string of source code or
    a byte sequence of raw bytecode.  For a module, it disassembles all functions.
    For a class, it disassembles all methods (including class and static methods).
    For a code object or sequence of raw bytecode, it prints one line per bytecode
-   instruction.  Strings are first compiled to code objects with the :func:`compile`
+   instruction.  It also recursively disassembles nested code objects (the code
+   of comprehensions, generator expressions and nested functions, and the code
+   used for building nested classes).
+   Strings are first compiled to code objects with the :func:`compile`
    built-in function before being disassembled.  If no object is provided, this
    function disassembles the last traceback.
 
    The disassembly is written as text to the supplied *file* argument if
    provided and to ``sys.stdout`` otherwise.
 
+   The maximal depth of recursion is limited by *depth* unless it is ``None``.
+   ``depth=0`` means no recursion.
+
    .. versionchanged:: 3.4
       Added *file* parameter.
 
+   .. versionchanged:: 3.7
+      Implemented recursive disassembling and added *depth* parameter.
+
 
 .. function:: distb(tb=None, *, file=None)
 
diff --git a/Doc/whatsnew/3.7.rst b/Doc/whatsnew/3.7.rst
index 6074781..3cdc009 100644
--- a/Doc/whatsnew/3.7.rst
+++ b/Doc/whatsnew/3.7.rst
@@ -178,6 +178,14 @@ contextlib
 :func:`contextlib.asynccontextmanager` has been added. (Contributed by
 Jelle Zijlstra in :issue:`29679`.)
 
+dis
+---
+
+The :func:`~dis.dis` function now is able to
+disassemble nested code objects (the code of comprehensions, generator
+expressions and nested functions, and the code used for building nested
+classes).  (Contributed by Serhiy Storchaka in :issue:`11822`.)
+
 distutils
 ---------
 
diff --git a/Lib/dis.py b/Lib/dis.py
index f3c18a5..b990839 100644
--- a/Lib/dis.py
+++ b/Lib/dis.py
@@ -31,7 +31,7 @@ def _try_compile(source, name):
         c = compile(source, name, 'exec')
     return c
 
-def dis(x=None, *, file=None):
+def dis(x=None, *, file=None, depth=None):
     """Disassemble classes, methods, functions, generators, or code.
 
     With no argument, disassemble the last traceback.
@@ -52,16 +52,16 @@ def dis(x=None, *, file=None):
             if isinstance(x1, _have_code):
                 print("Disassembly of %s:" % name, file=file)
                 try:
-                    dis(x1, file=file)
+                    dis(x1, file=file, depth=depth)
                 except TypeError as msg:
                     print("Sorry:", msg, file=file)
                 print(file=file)
     elif hasattr(x, 'co_code'): # Code object
-        disassemble(x, file=file)
+        _disassemble_recursive(x, file=file, depth=depth)
     elif isinstance(x, (bytes, bytearray)): # Raw bytecode
         _disassemble_bytes(x, file=file)
     elif isinstance(x, str):    # Source code
-        _disassemble_str(x, file=file)
+        _disassemble_str(x, file=file, depth=depth)
     else:
         raise TypeError("don't know how to disassemble %s objects" %
                         type(x).__name__)
@@ -338,6 +338,17 @@ def disassemble(co, lasti=-1, *, file=None):
     _disassemble_bytes(co.co_code, lasti, co.co_varnames, co.co_names,
                        co.co_consts, cell_names, linestarts, file=file)
 
+def _disassemble_recursive(co, *, file=None, depth=None):
+    disassemble(co, file=file)
+    if depth is None or depth > 0:
+        if depth is not None:
+            depth = depth - 1
+        for x in co.co_consts:
+            if hasattr(x, 'co_code'):
+                print(file=file)
+                print("Disassembly of %r:" % (x,), file=file)
+                _disassemble_recursive(x, file=file, depth=depth)
+
 def _disassemble_bytes(code, lasti=-1, varnames=None, names=None,
                        constants=None, cells=None, linestarts=None,
                        *, file=None, line_offset=0):
@@ -368,9 +379,9 @@ def _disassemble_bytes(code, lasti=-1, varnames=None, names=None,
         print(instr._disassemble(lineno_width, is_current_instr, offset_width),
               file=file)
 
-def _disassemble_str(source, *, file=None):
+def _disassemble_str(source, **kwargs):
     """Compile the source string, then disassemble the code object."""
-    disassemble(_try_compile(source, '<dis>'), file=file)
+    _disassemble_recursive(_try_compile(source, '<dis>'), **kwargs)
 
 disco = disassemble                     # XXX For backwards compatibility
 
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index e614b71..254b317 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -331,16 +331,77 @@ dis_fstring = """\
 def _g(x):
     yield x
 
+def _h(y):
+    def foo(x):
+        '''funcdoc'''
+        return [x + z for z in y]
+    return foo
+
+dis_nested_0 = """\
+%3d           0 LOAD_CLOSURE             0 (y)
+              2 BUILD_TUPLE              1
+              4 LOAD_CONST               1 (<code object foo at 0x..., file "%s", line %d>)
+              6 LOAD_CONST               2 ('_h.<locals>.foo')
+              8 MAKE_FUNCTION            8
+             10 STORE_FAST               1 (foo)
+
+%3d          12 LOAD_FAST                1 (foo)
+             14 RETURN_VALUE
+""" % (_h.__code__.co_firstlineno + 1,
+       __file__,
+       _h.__code__.co_firstlineno + 1,
+       _h.__code__.co_firstlineno + 4,
+)
+
+dis_nested_1 = """%s
+Disassembly of <code object foo at 0x..., file "%s", line %d>:
+%3d           0 LOAD_CLOSURE             0 (x)
+              2 BUILD_TUPLE              1
+              4 LOAD_CONST               1 (<code object <listcomp> at 0x..., file "%s", line %d>)
+              6 LOAD_CONST               2 ('_h.<locals>.foo.<locals>.<listcomp>')
+              8 MAKE_FUNCTION            8
+             10 LOAD_DEREF               1 (y)
+             12 GET_ITER
+             14 CALL_FUNCTION            1
+             16 RETURN_VALUE
+""" % (dis_nested_0,
+       __file__,
+       _h.__code__.co_firstlineno + 1,
+       _h.__code__.co_firstlineno + 3,
+       __file__,
+       _h.__code__.co_firstlineno + 3,
+)
+
+dis_nested_2 = """%s
+Disassembly of <code object <listcomp> at 0x..., file "%s", line %d>:
+%3d           0 BUILD_LIST               0
+              2 LOAD_FAST                0 (.0)
+        >>    4 FOR_ITER                12 (to 18)
+              6 STORE_FAST               1 (z)
+              8 LOAD_DEREF               0 (x)
+             10 LOAD_FAST                1 (z)
+             12 BINARY_ADD
+             14 LIST_APPEND              2
+             16 JUMP_ABSOLUTE            4
+        >>   18 RETURN_VALUE
+""" % (dis_nested_1,
+       __file__,
+       _h.__code__.co_firstlineno + 3,
+       _h.__code__.co_firstlineno + 3,
+)
+
 class DisTests(unittest.TestCase):
 
-    def get_disassembly(self, func, lasti=-1, wrapper=True):
+    maxDiff = None
+
+    def get_disassembly(self, func, lasti=-1, wrapper=True, **kwargs):
         # We want to test the default printing behaviour, not the file arg
         output = io.StringIO()
         with contextlib.redirect_stdout(output):
             if wrapper:
-                dis.dis(func)
+                dis.dis(func, **kwargs)
             else:
-                dis.disassemble(func, lasti)
+                dis.disassemble(func, lasti, **kwargs)
         return output.getvalue()
 
     def get_disassemble_as_string(self, func, lasti=-1):
@@ -350,7 +411,7 @@ class DisTests(unittest.TestCase):
         return re.sub(r'\b0x[0-9A-Fa-f]+\b', '0x...', text)
 
     def do_disassembly_test(self, func, expected):
-        got = self.get_disassembly(func)
+        got = self.get_disassembly(func, depth=0)
         if got != expected:
             got = self.strip_addresses(got)
         self.assertEqual(got, expected)
@@ -502,15 +563,29 @@ class DisTests(unittest.TestCase):
     def test_dis_object(self):
         self.assertRaises(TypeError, dis.dis, object())
 
+    def test_disassemble_recursive(self):
+        def check(expected, **kwargs):
+            dis = self.get_disassembly(_h, **kwargs)
+            dis = self.strip_addresses(dis)
+            self.assertEqual(dis, expected)
+
+        check(dis_nested_0, depth=0)
+        check(dis_nested_1, depth=1)
+        check(dis_nested_2, depth=2)
+        check(dis_nested_2, depth=3)
+        check(dis_nested_2, depth=None)
+        check(dis_nested_2)
+
+
 class DisWithFileTests(DisTests):
 
     # Run the tests again, using the file arg instead of print
-    def get_disassembly(self, func, lasti=-1, wrapper=True):
+    def get_disassembly(self, func, lasti=-1, wrapper=True, **kwargs):
         output = io.StringIO()
         if wrapper:
-            dis.dis(func, file=output)
+            dis.dis(func, file=output, **kwargs)
         else:
-            dis.disassemble(func, lasti, file=output)
+            dis.disassemble(func, lasti, file=output, **kwargs)
         return output.getvalue()
 
 
diff --git a/Misc/NEWS b/Misc/NEWS
index 8cbd463..b436524 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -355,6 +355,9 @@ Extension Modules
 Library
 -------
 
+- bpo-11822: The dis.dis() function now is able to disassemble nested
+  code objects.
+
 - bpo-30624: selectors does not take KeyboardInterrupt and SystemExit into
   account, leaving a fd in a bad state in case of error. Patch by Giampaolo
   Rodola'.