feat: change colorization logic(minimal set of colors, align with godbolt) & adjust test cases

Author: AbduazizZiyodov

Lib/_colorize.py

@@ -219,45 +219,18 @@ class Difflib(ThemeSection):
 
 @dataclass(frozen=True, kw_only=True)
 class Dis(ThemeSection):
-    label_bg: str = ANSIColors.BACKGROUND_CYAN
-    label_fg: str = ANSIColors.BLACK
+    disassembly_header: str = ANSIColors.GREEN
 
-    exception_label: str = ANSIColors.CYAN
-    argument_detail: str = "\x1B[3m"
+    jump_target: str = ANSIColors.GREEN
+    exception_label: str = ANSIColors.GREEN
 
-    op_load: str = ANSIColors.BLUE
-    op_pop: str = ANSIColors.MAGENTA
-    op_call_return: str = ANSIColors.YELLOW
-    op_control_flow: str = ANSIColors.GREEN
+    opname: str = ANSIColors.BLUE
+    opname_with_label: str = ANSIColors.GREEN
+
+    arg: str = ANSIColors.YELLOW
 
     reset: str = ANSIColors.RESET
 
-    def color_by_opname(self, opname: str) -> str:
-        if opname.startswith("LOAD_"):
-            return self.op_load
-
-        if opname.startswith("POP_"):
-            return self.op_pop
-
-        if opname.startswith(("CALL", "RETURN")) or opname in (
-            "YIELD_VALUE",
-            "MAKE_FUNCTION",
-            "SET_FUNCTION_ATTRIBUTE",
-            "RESUME",
-        ):
-            return self.op_call_return
-
-        if opname.startswith(("JUMP_", "POP_JUMP_", "FOR_ITER")) or opname in (
-            "SEND",
-            "GET_AWAITABLE",
-            "GET_AITER",
-            "GET_ANEXT",
-            "END_ASYNC_FOR",
-            "CLEANUP_THROW",
-        ):
-            return self.op_control_flow
-
-        return self.reset
 
 @dataclass(frozen=True, kw_only=True)
 class FancyCompleter(ThemeSection):

Lib/dis.py

@@ -525,7 +525,8 @@ def print_instruction_line(self, instr, mark_as_current):
         # Column: Label
         if instr.label is not None:
             lbl = f"L{instr.label}:"
-            fields.append(f"{lbl:>{label_width}}")
+            padded = f"{lbl:>{label_width}}"
+            fields.append(f"{theme.jump_target}{padded}{theme.reset}")
         else:
             fields.append(' ' * label_width)
         # Column: Instruction offset from start of code sequence
@@ -537,17 +538,18 @@ def print_instruction_line(self, instr, mark_as_current):
         else:
             fields.append('   ')
         # Column: Opcode name
-        fields.append(f"{theme.color_by_opname(instr.opname)}{instr.opname.ljust(_OPNAME_WIDTH)}{theme.reset}")
+        opname_color = theme.opname_with_label if instr.label is not None else theme.opname
+        fields.append(f"{opname_color}{instr.opname.ljust(_OPNAME_WIDTH)}{theme.reset}")
         # Column: Opcode argument
         if instr.arg is not None:
             # If opname is longer than _OPNAME_WIDTH, we allow it to overflow into
             # the space reserved for oparg. This results in fewer misaligned opargs
             # in the disassembly output.
             opname_excess = max(0, len(instr.opname) - _OPNAME_WIDTH)
-            fields.append(repr(instr.arg).rjust(_OPARG_WIDTH - opname_excess))
+            fields.append(f"{theme.arg}{repr(instr.arg)}{theme.reset}".rjust(_OPARG_WIDTH - opname_excess))
             # Column: Opcode argument details
             if instr.argrepr:
-                fields.append(f'{theme.argument_detail}(' + instr.argrepr + f'){theme.reset}')
+                fields.append('(' + instr.argrepr + ')')
         print(' '.join(fields).rstrip(), file=self.file)
 
     def print_exception_table(self, exception_entries):
@@ -591,6 +593,7 @@ def get_label_for_offset(self, offset):
         return self.labels_map.get(offset, None)
 
     def get_argval_argrepr(self, op, arg, offset):
+        theme = _get_dis_theme()
         get_name = None if self.names is None else self.names.__getitem__
         argval = None
         argrepr = ''
@@ -629,7 +632,7 @@ def get_argval_argrepr(self, op, arg, offset):
                 lbl = self.get_label_for_offset(argval)
                 assert lbl is not None
                 preposition = "from" if deop == END_ASYNC_FOR else "to"
-                argrepr = f"{preposition} L{lbl}"
+                argrepr = f"{preposition} {theme.jump_target}L{lbl}{theme.reset}"
             elif deop in (LOAD_FAST_LOAD_FAST, LOAD_FAST_BORROW_LOAD_FAST_BORROW, STORE_FAST_LOAD_FAST, STORE_FAST_STORE_FAST):
                 arg1 = arg >> 4
                 arg2 = arg & 15
@@ -870,7 +873,7 @@ def _disassemble_recursive(co, *, file=None, depth=None, show_caches=False, adap
         for x in co.co_consts:
             if hasattr(x, 'co_code'):
                 print(file=file)
-                print(f"{theme.label_bg}{theme.label_fg}Disassembly of {x!r}:{theme.reset}", file=file)
+                print(f"Disassembly of {theme.disassembly_header}{x!r}{theme.reset}:", file=file)
                 _disassemble_recursive(
                     x, file=file, depth=depth, show_caches=show_caches,
                     adaptive=adaptive, show_offsets=show_offsets,

Lib/test/test_dis.py

@@ -2680,6 +2680,7 @@ def test_specialized_code(self):
         for flag in ['-S', '--specialized']:
             self.check_output(source, expect, flag)
 
+
 @force_colorized_test_class
 class DisColoredTests(unittest.TestCase):
     def get_colored_output(self, func):
@@ -2690,55 +2691,132 @@ def get_colored_output(self, func):
 
         return output.getvalue()
 
-    def assertOpColored(self, output, opname, color):
-        self.assertIn(
-            f"{color}{opname}", output,
-            f"{opname} should be colored with {color!r}"
-        )
+    def _check_colored(self, output, opname, color, as_not_colored):
+        # allow spaces, ANSI colors etc.
+        inter_word_pattern = r"(?:\s|\x1b\[[0-9;]*m)*"
 
-    def test_load_ops_colored(self):
-        def f(a):
-            return a
-        out = self.get_colored_output(f)
-        self.assertOpColored(out, "LOAD_FAST", theme.op_load)
+        tokens = opname.split()
+        escaped_tokens = [re.escape(token) for token in tokens]
+        joined_opname = inter_word_pattern.join(escaped_tokens)
 
-    def test_call_return_ops_colored(self):
-        def f():
-            return 1
-        out = self.get_colored_output(f)
-        self.assertOpColored(out, "RETURN_VALUE", theme.op_call_return)
-        self.assertOpColored(out, "RESUME", theme.op_call_return)
+        pattern = re.escape(color) + inter_word_pattern + joined_opname
+
+        if as_not_colored:
+            self.assertNotRegex(
+                output,
+                pattern,
+                f"{opname} should NOT be colored with {color!r}",
+            )
+        else:
+            self.assertRegex(
+                output, pattern, f"{opname} should be colored with {color!r}"
+            )
 
-    def test_pop_ops_colored(self):
+    def assertOpColoredAs(self, output, opname, color):
+        self._check_colored(output, opname, color, as_not_colored=False)
+
+    def assertOpNotColoredAs(self, output, opname, wrong_color):
+        self._check_colored(output, opname, wrong_color, as_not_colored=True)
+
+    def test_opname_and_arg_colored(self):
         def f(a):
-            print(a)
+            return a
+
         out = self.get_colored_output(f)
-        self.assertOpColored(out, "POP_TOP", theme.op_pop)
+        self.assertOpColoredAs(out, "LOAD_FAST_BORROW", theme.opname)
+        self.assertOpColoredAs(out, "RETURN_VALUE", theme.opname)
+        self.assertOpColoredAs(out, "0", theme.arg)
 
     def test_control_flow_ops_colored(self):
         def f(a):
             for _ in a:
                 pass
+
         out = self.get_colored_output(f)
-        self.assertOpColored(out, "FOR_ITER", theme.op_control_flow)
-        self.assertOpColored(out, "JUMP_BACKWARD", theme.op_control_flow)
 
-    def test_argrepr_colored(self):
+        self.assertOpNotColoredAs(out, "FOR_ITER", theme.opname)
+        self.assertOpNotColoredAs(out, "END_FOR", theme.opname)
+
+        self.assertOpColoredAs(out, "FOR_ITER", theme.opname_with_label)
+        self.assertOpColoredAs(out, "END_FOR", theme.opname_with_label)
+
+        opnames = (
+            "RESUME",
+            "LOAD_FAST",
+            "GET_ITER",
+            "STORE_FAST",
+            "JUMP_BACKWARD",
+            "POP_ITER",
+            "LOAD_COMMON_CONSTANT",
+            "RETURN_VALUE",
+        )
+
+        for opname in opnames:
+            self.assertOpColoredAs(out, opname, theme.opname)
+
+    def test_jump_targets_colored(self):
+        # sample code from:
+        # https://github.com/python/cpython/pull/144208#issuecomment-5375286176
+        def f(a, c):
+            _t2.d if (
+                _t2 := (
+                    _t1
+                    if (_t1 := a.b if a is not None else None) is not None
+                    else c
+                )
+            ) is not None else None
+
+        out = self.get_colored_output(f)
+
+        for n in range(1, 6):
+            self.assertOpColoredAs(out, f"L{n}:", theme.jump_target)
+            self.assertIn(f"(to {theme.jump_target}L{n}{theme.reset})", out)
+
+        cases = (
+            "L1: LOAD_COMMON_CONSTANT",
+            "L2: COPY",
+            "L3: LOAD_FAST",
+            "L4: COPY",
+            "L5: LOAD_COMMON_CONSTANT",
+        )
+
+        for part in cases:
+            self.assertOpColoredAs(out, part, theme.jump_target)
+
+    def test_exception_table_colored(self):
         def f(a):
-            print(a)
+            try:
+                a
+            except Exception:
+                pass
+            else:
+                return a
+
         out = self.get_colored_output(f)
-        self.assertIn(f"{theme.argument_detail}(", out)
-
-    def test_color_by_opname_coverage(self):
-        self.assertEqual(theme.color_by_opname("LOAD_FAST"), theme.op_load)
-        self.assertEqual(theme.color_by_opname("LOAD_GLOBAL"), theme.op_load)
-        self.assertEqual(theme.color_by_opname("POP_TOP"), theme.op_pop)
-        self.assertEqual(theme.color_by_opname("CALL"), theme.op_call_return)
-        self.assertEqual(theme.color_by_opname("RETURN_VALUE"), theme.op_call_return)
-        self.assertEqual(theme.color_by_opname("RESUME"), theme.op_call_return)
-        self.assertEqual(theme.color_by_opname("FOR_ITER"), theme.op_control_flow)
-        self.assertEqual(theme.color_by_opname("JUMP_BACKWARD"), theme.op_control_flow)
-        self.assertEqual(theme.color_by_opname("BINARY_OP"), theme.reset)  # uncolored
+
+        cases = (
+            ("L1", "L2", "L3"),
+            ("L3", "L4", "L8"),
+            ("L5", "L6", "L8"),
+            ("L7", "L8", "L8"),
+        )
+
+        def assertExceptionTableRow(pairs, out):
+            p1, p2, p3 = pairs
+            part = f"{theme.jump_target}{p1}{theme.reset} to {theme.jump_target}{p2}{theme.reset} -> {theme.jump_target}{p3}{theme.reset}"
+            self.assertIn(part, out)
+
+        for pairs in cases:
+            assertExceptionTableRow(pairs, out)
+
+        cases = (
+            "L3: PUSH_EXC_INFO",
+            "L6: POP_EXCEPT",
+            "L7: RERAISE",
+        )
+
+        for part in cases:
+            self.assertOpColoredAs(out, part, theme.jump_target)
 
 if __name__ == "__main__":
     unittest.main()