More formatting

* Add 3.11+ COPY formatting
* Add ROT_THREE and ROT_FOUR
* expand use of NULL_EXTENDED_OP

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

@@ -60,7 +60,7 @@ def extended_format_binary_op(
                     arg1_start_offset, instructions, 1
                 )
                 if i is None:
-                    return "", None
+                    return NULL_EXTENDED_OP
             j = skip_cache(instructions, i + 1)
             stack_inst2 = instructions[j]
             if (
@@ -79,7 +79,7 @@ def extended_format_binary_op(
                 return fmt_str % (arg2, arg1), start_offset
             else:
                 return fmt_str % ("...", arg1), None
-    return "", None
+    return NULL_EXTENDED_OP
 
 
 def extended_format_infix_binary_op(
@@ -103,7 +103,7 @@ def extended_format_infix_binary_op(
         if arg1_start_offset is not None:
             i = get_instruction_index_from_offset(arg1_start_offset, instructions, 1)
             if i is None:
-                return "", None
+                return NULL_EXTENDED_OP
         j = i + 1
         # 3.11+ has CACHE instructions
         while instructions[j].opname == "CACHE":
@@ -129,7 +129,7 @@ def extended_format_infix_binary_op(
             return f"{arg2}{op_str}{arg1}", start_offset
         else:
             return f"...{op_str}{arg1}", None
-    return "", None
+    return NULL_EXTENDED_OP
 
 
 def extended_format_store_op(
@@ -142,7 +142,7 @@ def extended_format_store_op(
     # are more complicated, so let's not try to figure this out.
     # This kind of things is best left for a decompiler.
     if inst.is_jump_target:
-        return "", None
+        return NULL_EXTENDED_OP
 
     prev_inst = instructions[1]
     start_offset = prev_inst.offset
@@ -198,7 +198,7 @@ def extended_format_ternary_op(
         if arg1_start_offset is not None:
             i = get_instruction_index_from_offset(arg1_start_offset, instructions, 1)
             if i is None:
-                return "", None
+                return NULL_EXTENDED_OP
         j = skip_cache(instructions, i + 1)
         stack_inst2 = instructions[j]
         if (
@@ -229,7 +229,7 @@ def extended_format_ternary_op(
             return fmt_str % (arg2, arg1, arg3), start_offset
         else:
             return fmt_str % ("...", "...", "..."), None
-    return "", None
+    return NULL_EXTENDED_OP
 
 
 def extended_format_STORE_SUBSCR(
@@ -251,7 +251,7 @@ def extended_format_unary_op(
         return fmt_str % stack_arg.tos_str, start_offset
     if stack_arg.opcode in opc.operator_set:
         return fmt_str % stack_arg.argrepr, start_offset
-    return "", None
+    return NULL_EXTENDED_OP
 
 
 def extended_format_ATTR(
@@ -271,7 +271,7 @@ def extended_format_ATTR(
             f"{base}.{instructions[0].argrepr}",
             instr1.start_offset,
         )
-    return "", None
+    return NULL_EXTENDED_OP
 
 
 def extended_format_BINARY_ADD(
@@ -372,7 +372,7 @@ def extended_format_build_tuple_or_list(
             return f"{left_delim}{args_str},{right_delim}", instructions[i].start_offset
         else:
             return f"{left_delim}{args_str}{right_delim}", instructions[i].start_offset
-    return "", None
+    return NULL_EXTENDED_OP
 
 
 def extended_format_BUILD_CONST_KEY_MAP(opc, instructions):
@@ -393,7 +393,7 @@ def extended_format_BUILD_CONST_KEY_MAP(opc, instructions):
                 arg_pairs.append(f"{key_values[i]}: {arglist[i]}")
             args_str = ", ".join(arg_pairs)
             return "{" + args_str + "}", instructions[i].start_offset
-    return "", None
+    return NULL_EXTENDED_OP
 
 
 def extended_format_BUILD_LIST(
@@ -426,7 +426,7 @@ def extended_format_BUILD_SLICE(
     if instructions[0].argval == 0:
         # Degenerate case
         return "set()", instructions[0].start_offset
-    return "", None
+    return NULL_EXTENDED_OP
 
 
 def extended_format_BUILD_TUPLE(
@@ -479,7 +479,7 @@ def extended_format_CALL_FUNCTION(opc, instructions) -> Tuple[str, Optional[int]
     arglist, arg_count, i = get_arglist(instructions, 0, arg_count)
 
     if arglist is None:
-        return "", None
+        return NULL_EXTENDED_OP
 
     assert i is not None
     if i >= len(instructions) - 1:
@@ -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,
     },
 }