Add 3.11+ COPY formatting

Author: rocky

xdis/opcodes/format/basic.py

@@ -58,6 +58,14 @@ def format_RAISE_VARARGS_older(argc) -> str:
         return "exception, parameter, traceback"
     return ""
 
+def format_ROT_FOUR(_: int) -> str:
+    return "TOS, TOS1, TOS2, TOS3 = TOS1, TOS2, TOS3, TOS"
+
+
+def format_ROT_THREE(_: int) -> str:
+    return "TOS, TOS1, TOS2 = TOS1, TOS2, TOS"
+
+
 def format_ROT_TWO(_: int) -> str:
     # We add a space at the end as a sentinal to use in get_instruction_tos_str()
     return "TOS, TOS1 = TOS1, TOS"
@@ -70,5 +78,7 @@ def format_ROT_TWO(_: int) -> str:
     "CALL_FUNCTION_VAR_KW": format_CALL_FUNCTION_pos_name_encoded,
     "EXTENDED_ARG": format_extended_arg,
     "RAISE_VARARGS": format_RAISE_VARARGS_older,
+    "ROT_FOUR": format_ROT_FOUR,
+    "ROT_THREE": format_ROT_THREE,
     "ROT_TWO": format_ROT_TWO,
 }

xdis/opcodes/format/extended.py

@@ -791,9 +791,9 @@ def get_instruction_arg(inst: Instruction, argval=None) -> str:
 def get_instruction_tos_str(inst: Instruction) -> str:
     if inst.tos_str is not None:
         argval = inst.tos_str
-        argval_without_push = re.match(r"^push\((.+)\) ", argval)
+        argval_without_push = re.match(r"^(?:push|copy)\((.+)\) ", argval)
         if argval_without_push:
-            # remove surrounding "push(...)" string
+            # remove surrounding "push(...) or copy(...)" string
             argval = argval_without_push.group(1)
     else:
         argval = inst.argrepr

xdis/opcodes/opcode_311.py

@@ -36,7 +36,11 @@
     store_op,
     update_pj3,
 )
-from xdis.opcodes.format.extended import extended_format_binary_op
+from xdis.opcodes.format.extended import (
+    NULL_EXTENDED_OP,
+    extended_format_binary_op,
+    extended_format_unary_op,
+)
 from xdis.opcodes.opcode_310 import opcode_arg_fmt310, opcode_extended_fmt310
 
 version_tuple = (3, 11)
@@ -245,6 +249,21 @@ def extended_format_BINARY_OP(opc, instructions) -> Tuple[str, Optional[int]]:
     return extended_format_binary_op(opc, instructions, f"%s {opname} %s")
 
 
+def extended_format_COPY_OP(
+    opc, instructions: List[Instruction]
+) -> Tuple[str, Optional[int]]:
+    """Try to extract TOS value and show that surrounded in a "push() ".
+    The trailing space at the used as a sentinal for `get_instruction_tos_str()`
+    which tries to remove the push() part when the operand value string is needed.
+    """
+
+    # We add a space at the end as a sentinal to use in get_instruction_tos_str()
+    if instructions[1].optype not in ["jrel", "jabs"]:
+        return extended_format_unary_op(opc, instructions, "copy(%s) ")
+    else:
+        return NULL_EXTENDED_OP
+
+
 def extended_format_SWAP(
     opc, instructions: List[Instruction]
 ) -> Tuple[str, Optional[int]]:
@@ -263,12 +282,13 @@ def extended_format_SWAP(
     i = swap_instr.argval
     # s = ""
 
-    if (i is None or not (0 < i < len(instructions))):
+    if i is None or not (0 < i < len(instructions)):
         return "", None
 
     # To be continued
     return "", None
 
+
 def format_BINARY_OP(arg: int) -> str:
     return _nb_ops[arg][1]
 
@@ -294,6 +314,7 @@ def format_SWAP_OP(arg: int) -> str:
     **opcode_extended_fmt310,
     **{
         "BINARY_OP": extended_format_BINARY_OP,
+        "COPY": extended_format_COPY_OP,
     },
 }