gh-112720: make it easier to subclass and modify dis.ArgResolver's jump arg resolution (#115564)

2 files changed, 35 insertions, 11 deletions

diff --git a/Lib/dis.py b/Lib/dis.py
index f05ea1a..d146bcb 100644
--- a/Lib/dis.py
+++ b/Lib/dis.py
@@ -505,6 +505,20 @@ class ArgResolver:
         self.varname_from_oparg = varname_from_oparg
         self.labels_map = labels_map or {}
 
+    def offset_from_jump_arg(self, op, arg, offset):
+        deop = _deoptop(op)
+        if deop in hasjabs:
+            return arg * 2
+        elif deop in hasjrel:
+            signed_arg = -arg if _is_backward_jump(deop) else arg
+            argval = offset + 2 + signed_arg*2
+            caches = _get_cache_size(_all_opname[deop])
+            argval += 2 * caches
+            if deop == ENTER_EXECUTOR:
+                argval += 2
+            return argval
+        return None
+
     def get_label_for_offset(self, offset):
         return self.labels_map.get(offset, None)
 
@@ -536,17 +550,11 @@ class ArgResolver:
                         argrepr = f"{argrepr} + NULL|self"
                 else:
                     argval, argrepr = _get_name_info(arg, get_name)
-            elif deop in hasjabs:
-                argval = arg*2
-                argrepr = f"to L{self.labels_map[argval]}"
-            elif deop in hasjrel:
-                signed_arg = -arg if _is_backward_jump(deop) else arg
-                argval = offset + 2 + signed_arg*2
-                caches = _get_cache_size(_all_opname[deop])
-                argval += 2 * caches
-                if deop == ENTER_EXECUTOR:
-                    argval += 2
-                argrepr = f"to L{self.labels_map[argval]}"
+            elif deop in hasjump or deop in hasexc:
+                argval = self.offset_from_jump_arg(op, arg, offset)
+                lbl = self.get_label_for_offset(argval)
+                assert lbl is not None
+                argrepr = f"to L{lbl}"
             elif deop in (LOAD_FAST_LOAD_FAST, STORE_FAST_LOAD_FAST, STORE_FAST_STORE_FAST):
                 arg1 = arg >> 4
                 arg2 = arg & 15
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index a5917da..a93cb50 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -1986,6 +1986,22 @@ class InstructionTests(InstructionTestCase):
         self.assertEqual(f(opcode.opmap["BINARY_OP"], 3, *args), (3, '<<'))
         self.assertEqual(f(opcode.opmap["CALL_INTRINSIC_1"], 2, *args), (2, 'INTRINSIC_IMPORT_STAR'))
 
+    def test_custom_arg_resolver(self):
+        class MyArgResolver(dis.ArgResolver):
+            def offset_from_jump_arg(self, op, arg, offset):
+                return arg + 1
+
+            def get_label_for_offset(self, offset):
+                return 2 * offset
+
+        def f(opcode, oparg, offset, *init_args):
+            arg_resolver = MyArgResolver(*init_args)
+            return arg_resolver.get_argval_argrepr(opcode, oparg, offset)
+        offset = 42
+        self.assertEqual(f(opcode.opmap["JUMP_BACKWARD"], 1, offset), (2, 'to L4'))
+        self.assertEqual(f(opcode.opmap["SETUP_FINALLY"], 2, offset), (3, 'to L6'))
+
+
     def get_instructions(self, code):
         return dis._get_instructions_bytes(code)