Add --method option to limit to specific methods

Author: rocky

xdis/bin/pydisasm.py

@@ -37,14 +37,22 @@
     ),
     help="Select disassembly style",
 )
+@click.option(
+    "--method",
+    "-m",
+    multiple=True,
+    type=str,
+    help="Specify a method name to filter disassembly by. Can be given multiple times.",
+)
+
 @click.option(
     "--show-source/--no-show-source",
     "-S",
     help="Intersperse Python source text from linecache if available.",
 )
 @click.version_option(version=__version__)
 @click.argument("files", nargs=-1, type=click.Path(readable=True), required=True)
-def main(format, show_source: bool, files):
+def main(format, method: tuple, show_source: bool, files):
     """Disassembles a Python bytecode file.
 
     We handle bytecode for virtually every release of Python and some releases of PyPy.
@@ -80,7 +88,7 @@ def main(format, show_source: bool, files):
             continue
 
         try:
-            disassemble_file(path, sys.stdout, format, show_source=show_source)
+            disassemble_file(path, sys.stdout, format, show_source=show_source, methods=method)
         except ImportError as e:
             print(e)
             rc = 3

xdis/disasm.py

@@ -29,6 +29,7 @@
 import sys
 import types
 from collections import deque
+from typing import Tuple
 
 import xdis
 from xdis.bytecode import Bytecode
@@ -135,6 +136,7 @@ def disco(
     alternate_opmap=None,
     show_source=False,
     is_graal=False,
+    methods=tuple(),
 ):
     """
     disassembles and deparses a given code block 'co'
@@ -160,7 +162,8 @@ def disco(
     real_out = out or sys.stdout
 
     if co.co_filename and asm_format != "xasm":
-        real_out.write(format_code_info(co, version_tuple, is_graal=is_graal) + "\n")
+        if not_filtered(co, methods):
+            real_out.write(format_code_info(co, version_tuple, is_graal=is_graal) + "\n")
         pass
 
     opc = get_opcode(version_tuple, is_pypy, alternate_opmap)
@@ -180,6 +183,7 @@ def disco(
             asm_format=asm_format,
             dup_lines=True,
             show_source=show_source,
+            methods=methods,
         )
 
 
@@ -191,6 +195,7 @@ def disco_loop(
     dup_lines=False,
     asm_format="classic",
     show_source=False,
+    methods=tuple(),
 ):
     """Disassembles a queue of code objects. If we discover
     another code object which will be found in co_consts, we add
@@ -204,25 +209,26 @@ def disco_loop(
 
     while len(queue) > 0:
         co = queue.popleft()
-        if co.co_name not in ("<module>", "?"):
-            real_out.write("\n" + format_code_info(co, version_tuple) + "\n")
-
-        if asm_format == "dis":
-            assert version_tuple[:2] == PYTHON_VERSION_TRIPLE[:2], (
-                "dis requires disassembly from the same Python version: "
-                f"Bytecode is for {version_tuple[:2]}; Running:{PYTHON_VERSION_TRIPLE[:2]}"
-            )
-            dis.disassemble(co, lasti=-1, file=real_out)
-        else:
-            bytecode = Bytecode(co, opc, dup_lines=dup_lines)
-            real_out.write(
-                bytecode.dis(asm_format=asm_format, show_source=show_source) + "\n"
-            )
+        if not_filtered(co, methods):
+            if co.co_name not in ("<module>", "?"):
+                real_out.write("\n" + format_code_info(co, version_tuple) + "\n")
+
+            if asm_format == "dis":
+                assert version_tuple[:2] == PYTHON_VERSION_TRIPLE[:2], (
+                    "dis requires disassembly from the same Python version: "
+                    f"Bytecode is for {version_tuple[:2]}; Running:{PYTHON_VERSION_TRIPLE[:2]}"
+                )
+                dis.disassemble(co, lasti=-1, file=real_out)
+            else:
+                bytecode = Bytecode(co, opc, dup_lines=dup_lines)
+                real_out.write(
+                    bytecode.dis(asm_format=asm_format, show_source=show_source) + "\n"
+                )
 
-            if version_tuple >= (3, 11):
-                if bytecode.exception_entries not in (None, []):
-                    exception_table = format_exception_table(bytecode, version_tuple)
-                    real_out.write(exception_table + "\n")
+                if version_tuple >= (3, 11):
+                    if bytecode.exception_entries not in (None, []):
+                        exception_table = format_exception_table(bytecode, version_tuple)
+                        real_out.write(exception_table + "\n")
 
         for c in co.co_consts:
             if iscode(c):
@@ -312,6 +318,7 @@ def disassemble_file(
     asm_format="classic",
     alternate_opmap=None,
     show_source=False,
+    methods: Tuple[str] = tuple()
 ):
     """
     Disassemble Python byte-code file (.pyc).
@@ -383,6 +390,7 @@ def disassemble_file(
             alternate_opmap=alternate_opmap,
             show_source=show_source,
             is_graal=is_graal,
+            methods=methods,
         )
     # print co.co_filename
     return (
@@ -396,13 +404,16 @@ def disassemble_file(
         sip_hash,
     )
 
+def not_filtered(co: types.CodeType, methods: tuple) -> bool:
+    return len(methods) == 0 or co.co_name in methods
 
 def _test():
     """Simple test program to disassemble a file."""
     argc = len(sys.argv)
     if argc == 1:
         if xdis.PYTHON3:
             disassemble_file(__file__)
+            disassemble_file(__file__, methods=("code_uniquify",))
         else:
             sys.stderr.write(f"usage: {__file__} [-|CPython compiled file [format]]\n")
             sys.exit(2)

xdis/marsh.py

@@ -29,15 +29,11 @@
 
 import struct
 import types
+from sys import intern
 
 from xdis.codetype import Code2, Code3
 from xdis.version_info import PYTHON3, PYTHON_VERSION_TRIPLE, version_tuple_to_str
 
-try:
-    intern
-except NameError:
-    from sys import intern
-
 try:
     from __pypy__ import builtinify