Issue #21947: handle generator-iterator objects in dis

Patch by Clement Rouault.

2 files changed, 14 insertions, 2 deletions

diff --git a/Lib/dis.py b/Lib/dis.py
index 81cbe7f4f9..d215bc59e4 100644
--- a/Lib/dis.py
+++ b/Lib/dis.py
@@ -29,7 +29,7 @@ def _try_compile(source, name):
     return c
 
 def dis(x=None, *, file=None):
-    """Disassemble classes, methods, functions, or code.
+    """Disassemble classes, methods, functions, generators, or code.
 
     With no argument, disassemble the last traceback.
 
@@ -41,6 +41,8 @@ def dis(x=None, *, file=None):
         x = x.__func__
     if hasattr(x, '__code__'):  # Function
         x = x.__code__
+    if hasattr(x, 'gi_code'):  # Generator
+        x = x.gi_code
     if hasattr(x, '__dict__'):  # Class or module
         items = sorted(x.__dict__.items())
         for name, x1 in items:
@@ -99,11 +101,13 @@ def pretty_flags(flags):
     return ", ".join(names)
 
 def _get_code_object(x):
-    """Helper to handle methods, functions, strings and raw code objects"""
+    """Helper to handle methods, functions, generators, strings and raw code objects"""
     if hasattr(x, '__func__'): # Method
         x = x.__func__
     if hasattr(x, '__code__'): # Function
         x = x.__code__
+    if hasattr(x, 'gi_code'):  # Generator
+        x = x.gi_code
     if isinstance(x, str):     # Source code
         x = _try_compile(x, "<disassembly>")
     if hasattr(x, 'co_code'):  # Code object
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index d1229fba6d..f04f12ca82 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -229,6 +229,9 @@ dis_traceback = """\
        TRACEBACK_CODE.co_firstlineno + 4,
        TRACEBACK_CODE.co_firstlineno + 5)
 
+def _g(x):
+    yield x
+
 class DisTests(unittest.TestCase):
 
     def get_disassembly(self, func, lasti=-1, wrapper=True):
@@ -314,6 +317,11 @@ class DisTests(unittest.TestCase):
         method_bytecode = _C(1).__init__.__code__.co_code
         self.do_disassembly_test(method_bytecode, dis_c_instance_method_bytes)
 
+    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
+        self.assertEqual(gen_disas, gen_func_disas)
+
     def test_dis_none(self):
         try:
             del sys.last_traceback