bpo-31183: `dis` now handles coroutines & async generators (GH-3077)

Coroutines and async generators use a distinct attribute name for their
code objects, so this updates the `dis` module to correctly disassemble
objects with those attributes.

Due to the increase in the test module length, it also fixes some latent
defects in the tests related to how the displayed source line numbers
are extracted.

https://bugs.python.org/issue31230 is a follow-up issue suggesting we
may want to solve this a different way, by instead giving all these object
types a common `__code__` attribute, avoiding the need for special
casing in the `dis` module.

1 files changed, 27 insertions, 5 deletions

diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index 254b317e49..bfbbee2794 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -331,6 +331,13 @@ dis_fstring = """\
 def _g(x):
     yield x
 
+async def _ag(x):
+    yield x
+
+async def _co(x):
+    async for item in _ag(x):
+        pass
+
 def _h(y):
     def foo(x):
         '''funcdoc'''
@@ -390,6 +397,7 @@ Disassembly of <code object <listcomp> at 0x..., file "%s", line %d>:
        _h.__code__.co_firstlineno + 3,
 )
 
+
 class DisTests(unittest.TestCase):
 
     maxDiff = None
@@ -531,10 +539,22 @@ class DisTests(unittest.TestCase):
         self.do_disassembly_test(_C.cm, dis_c_class_method)
 
     def test_disassemble_generator(self):
-        gen_func_disas = self.get_disassembly(_g)  # Disassemble generator function
-        gen_disas = self.get_disassembly(_g(1))  # Disassemble generator itself
+        gen_func_disas = self.get_disassembly(_g)  # Generator function
+        gen_disas = self.get_disassembly(_g(1))  # Generator iterator
         self.assertEqual(gen_disas, gen_func_disas)
 
+    def test_disassemble_async_generator(self):
+        agen_func_disas = self.get_disassembly(_ag)  # Async generator function
+        agen_disas = self.get_disassembly(_ag(1))  # Async generator iterator
+        self.assertEqual(agen_disas, agen_func_disas)
+
+    def test_disassemble_coroutine(self):
+        coro_func_disas = self.get_disassembly(_co)  # Coroutine function
+        coro = _co(1)  # Coroutine object
+        coro.close()  # Avoid a RuntimeWarning (never awaited)
+        coro_disas = self.get_disassembly(coro)
+        self.assertEqual(coro_disas, coro_func_disas)
+
     def test_disassemble_fstring(self):
         self.do_disassembly_test(_fstring, dis_fstring)
 
@@ -1051,11 +1071,13 @@ class BytecodeTests(unittest.TestCase):
 
     def test_source_line_in_disassembly(self):
         # Use the line in the source code
-        actual = dis.Bytecode(simple).dis()[:3]
-        expected = "{:>3}".format(simple.__code__.co_firstlineno)
+        actual = dis.Bytecode(simple).dis()
+        actual = actual.strip().partition(" ")[0]  # extract the line no
+        expected = str(simple.__code__.co_firstlineno)
         self.assertEqual(actual, expected)
         # Use an explicit first line number
-        actual = dis.Bytecode(simple, first_line=350).dis()[:3]
+        actual = dis.Bytecode(simple, first_line=350).dis()
+        actual = actual.strip().partition(" ")[0]  # extract the line no
         self.assertEqual(actual, "350")
 
     def test_info(self):