Improve disassembly output of compiler --dis

Summary: Currently when you use --dis we just use the built-in disassembler. This pulls our test disassembler out and re-uses it so you can see the extended opcodes in a sensible way:

Reviewed By: yoney

Differential Revision: D103307684

fbshipit-source-id: abc1b2d3f22150adb0246a3dbee1bc3ae532b49a

Author: DinoV

cinderx/PythonLib/cinderx/compiler/__init__.py

@@ -15,12 +15,61 @@
 """
 
 import ast
+import dis
+import sys
+from io import StringIO
 from types import CodeType
 from typing import Any
 
+from .opcodes import STATIC_CONST_OPCODES, STATIC_OPNAMES
 from .pycodegen import CinderCodeGenerator, compile, compile_code, compileFile
 
 
+def make_static_instr(instr: dis.Instruction, co: object) -> dis.Instruction:
+    if instr.opcode in STATIC_CONST_OPCODES:
+        # pyre-fixme[21]: Could not find name `_get_code_object` in `dis` (stubbed).
+        from dis import _get_code_object
+
+        return dis.Instruction(
+            STATIC_OPNAMES[instr.opcode],
+            instr[1],
+            instr.arg,
+            _get_code_object(co).co_consts[instr.arg],
+            *instr[4:],
+        )
+    return dis.Instruction(STATIC_OPNAMES[instr.opcode], *instr[1:])
+
+
+def get_disassembly_as_string(co: object, recurse: bool = False) -> str:
+    s = StringIO()
+    if sys.version_info < (3, 14):
+        dis.dis(co, file=s)
+        return s.getvalue()
+
+    # pyre-fixme[21]: Could not find name `Formatter` in `dis` (stubbed).
+    # pyre-fixme[21]: Could not find name `_get_code_object` in `dis` (stubbed).
+    from dis import _get_code_object, Bytecode, Formatter
+
+    formatter = Formatter(file=s, offset_width=3)
+    bc = Bytecode(co)
+    extended = False
+    for instr in bc:
+        if extended and instr.opname != "EXTENDED_ARG":
+            extended = False
+            instr = make_static_instr(instr, co)
+        elif instr.opname == "EXTENDED_OPCODE":
+            extended = True
+        formatter.print_instruction(instr, False)
+
+    if recurse:
+        for const in _get_code_object(co).co_consts:
+            if isinstance(const, CodeType):
+                s.write(f"\nDisassembly of {const!r}:\n")
+                s.write(get_disassembly_as_string(const, recurse=True))
+
+    return s.getvalue()
+
+
 def exec_cinder(
     source: str | bytes | ast.Module | ast.Expression | ast.Interactive | CodeType,
     locals: dict[str, Any],
@@ -37,4 +86,11 @@ def exec_cinder(
     exec(code, locals, globals)
 
 
-__all__ = ("compile", "compile_code", "compileFile", "exec_cinder")
+__all__ = (
+    "compile",
+    "compile_code",
+    "compileFile",
+    "exec_cinder",
+    "get_disassembly_as_string",
+    "make_static_instr",
+)

cinderx/PythonLib/cinderx/compiler/__main__.py

@@ -16,6 +16,7 @@
 from types import CodeType
 from typing import Pattern, TextIO
 
+from . import get_disassembly_as_string
 from .pycodegen import CinderCodeGenerator, compile_code, make_header
 from .static import FIXED_MODULES, StaticCodeGenerator
 from .static.pyrefly_info import Pyrefly
@@ -236,7 +237,7 @@ def main() -> None:
             )
 
         if args.dis:
-            dis(codeobj)
+            print(get_disassembly_as_string(codeobj, recurse=True), end="")
 
         if args.c:
             if args.output:

cinderx/PythonLib/test_cinderx/test_compiler/common.py

@@ -66,42 +66,16 @@ class CompilerTest(TestCase):
     COMPARE_JUMP_ZERO = "POP_JUMP_IF_ZERO"
 
     def get_disassembly_as_string(self, co: Disassembleable) -> str:
-        s = StringIO()
-        if sys.version_info < (3, 14):
-            dis.dis(co, file=s)
-            return s.getvalue()
-
-        # pyre-fixme[10]: Name `Formatter` is used but not defined.
-        # pyre-fixme[16]: Module `dis` has no attribute `Formatter`.
-        formatter = Formatter(file=s, offset_width=3)
-        # pyre-fixme[10]: Name `Bytecode` is used but not defined.
-        bc = Bytecode(co)
-        extended = False
-        for instr in bc:
-            if extended and instr.opname != "EXTENDED_ARG":
-                extended = False
-                instr = self.make_static_instr(instr, co)
-            elif instr.opname == "EXTENDED_OPCODE":
-                extended = True
+        from cinderx.compiler import get_disassembly_as_string
 
-            formatter.print_instruction(instr, False)
-
-        return s.getvalue()
+        return get_disassembly_as_string(co)
 
     def make_static_instr(
         self, instr: dis.Instruction, co: Disassembleable
     ) -> dis.Instruction:
-        if instr.opcode in STATIC_CONST_OPCODES:
-            return dis.Instruction(
-                STATIC_OPNAMES[instr.opcode],
-                instr[1],
-                instr.arg,
-                # pyre-fixme[10]: Name `_get_code_object` is used but not defined.
-                # pyre-fixme[16]: Module `dis` has no attribute `_get_code_object`.
-                _get_code_object(co).co_consts[instr.arg],
-                *instr[4:],
-            )
-        return dis.Instruction(STATIC_OPNAMES[instr.opcode], *instr[1:])
+        from cinderx.compiler import make_static_instr
+
+        return make_static_instr(instr, co)
 
     def assertInBytecode(
         self,