From 7f75ef1c5572803614e14664f82e6986338254a3 Mon Sep 17 00:00:00 2001
From: Shahar Naveh <50263213+ShaharNaveh@users.noreply.github.com>
Date: Wed, 10 Jun 2026 05:17:50 +0300
Subject: [PATCH] `scripts/update_lib migrate` to preserve patches on classes
 (#8057)

---
 scripts/update_lib/patch_spec.py            | 270 +++++++++++++-------
 scripts/update_lib/tests/test_patch_spec.py |  26 ++
 2 files changed, 198 insertions(+), 98 deletions(-)

diff --git a/scripts/update_lib/patch_spec.py b/scripts/update_lib/patch_spec.py
index 565d01e3e..996b78bfa 100644
--- a/scripts/update_lib/patch_spec.py
+++ b/scripts/update_lib/patch_spec.py
@@ -6,9 +6,13 @@ This module handles:
 - Applying patches to test files (JSON -> file)
 """
 
+from __future__ import annotations
+
 import ast
 import collections
+import contextlib
 import enum
+import re
 import textwrap
 import typing
 
@@ -91,33 +95,109 @@ class PatchSpec(typing.NamedTuple):
 
         return f"@{unparsed}"
 
+    @classmethod
+    def try_from_ast_node(
+        cls, node: ast.Attribute | ast.Call, lines: list[str]
+    ) -> typing.Self | None:
+        if isinstance(node, ast.Attribute):
+            attr_node = node
+        elif isinstance(node, ast.Call):
+            attr_node = node.func
+        else:
+            return
 
-def _single_to_double_quotes(s: str) -> str:
-    """Convert single-quoted strings to double-quoted strings.
+        if (
+            isinstance(attr_node, ast.Name)
+            or getattr(attr_node.value, "id", None) != UT
+        ):
+            return
 
-    Falls back to original if conversion breaks the AST equivalence.
-    """
-    import re
+        cond = None
+        try:
+            ut_method = UtMethod(attr_node.attr)
+        except ValueError:
+            return
 
-    def replace_string(match: re.Match) -> str:
-        content = match.group(1)
-        # Unescape single quotes and escape double quotes
-        content = content.replace("\\'", "'").replace('"', '\\"')
-        return f'"{content}"'
+        # If our ut_method has args then,
+        # we need to search for a constant that contains our `COMMENT`.
+        # Otherwise we need to search it in the raw source code :/
+        if ut_method.has_args():
+            reason = next(
+                (
+                    inner_node.value
+                    for inner_node in ast.walk(node)
+                    if isinstance(inner_node, ast.Constant)
+                    and isinstance(inner_node.value, str)
+                    and COMMENT in inner_node.value
+                ),
+                None,
+            )
 
-    # Match single-quoted strings (handles escaped single quotes inside)
-    converted = re.sub(r"'((?:[^'\\]|\\.)*)'", replace_string, s)
+            # If we didn't find a constant containing <COMMENT>,
+            # then we didn't put this decorator
+            if not reason:
+                return
 
-    # Verify: parse converted and unparse should equal original
-    try:
-        converted_ast = ast.parse(converted, mode="eval")
-        if ast.unparse(converted_ast) == s:
-            return converted
-    except SyntaxError:
-        pass
+            if ut_method.has_cond():
+                cond = ast.unparse(node.args[0])
+        else:
+            pattern = re.compile(rf"{COMMENT}.?(.*)")
+            dec_lineno = node.lineno
 
-    # Fall back to original if conversion failed
-    return s
+            curr_line = lines[dec_lineno - 1]
+            prev_line = lines[dec_lineno - 2]
+
+            # If we see our comment at the decorator line, take it
+            if found := pattern.search(curr_line):
+                reason = found.group()
+            elif prev_line.strip().startswith("#") and (
+                found := pattern.search(prev_line)
+            ):
+                # Search the previous line of the decorator,
+                # only take the comment if the line starts with a `#`
+                reason = found.group()
+            else:
+                # Didn't find our `COMMENT`, so the patch isn't ours :)
+                return
+
+        reason = reason.removeprefix(COMMENT).strip(";:, ")
+        return cls(ut_method, cond, reason)
+
+
+class PatchEntryVisitor(ast.NodeVisitor):
+    def __init__(self, lines: list[str]):
+        self.current_class = None
+        self.patches = []
+        self.lines = lines
+
+    def patches_from_node(
+        self, node: ast.FunctionDef | ast.AsyncFunctionDef
+    ) -> Iterator[PatchEntry]:
+        for dec_node in node.decorator_list:
+            spec = PatchSpec.try_from_ast_node(dec_node, self.lines)
+
+            if spec is None:
+                continue
+
+            yield PatchEntry(self.current_class, node.name, spec)
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
+        self.patches.extend(self.patches_from_node(node))
+        # TODO: Support nested classes/methods
+        # self.generic_visit(node)
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        self.patches.extend(self.patches_from_node(node))
+        # TODO: Support nested classes/methods
+        # self.generic_visit(node)
+
+    def visit_ClassDef(self, node: ast.ClassDef):
+        with temp_attr(self, "current_class", node.name):
+            for patch in self.patches_from_node(node):
+                patch = patch._replace(test_name="__self__")
+                self.patches.append(patch)
+
+            self.generic_visit(node)
 
 
 class PatchEntry(typing.NamedTuple):
@@ -142,76 +222,10 @@ class PatchEntry(typing.NamedTuple):
     def iter_patch_entries(
         cls, tree: ast.Module, lines: list[str]
     ) -> "Iterator[typing.Self]":
-        import re
-        import sys
 
-        for cls_node, fn_node in iter_tests(tree):
-            parent_class = cls_node.name
-            for dec_node in fn_node.decorator_list:
-                if not isinstance(dec_node, (ast.Attribute, ast.Call)):
-                    continue
-
-                attr_node = (
-                    dec_node if isinstance(dec_node, ast.Attribute) else dec_node.func
-                )
-
-                if (
-                    isinstance(attr_node, ast.Name)
-                    or getattr(attr_node.value, "id", None) != UT
-                ):
-                    continue
-
-                cond = None
-                try:
-                    ut_method = UtMethod(attr_node.attr)
-                except ValueError:
-                    continue
-
-                # If our ut_method has args then,
-                # we need to search for a constant that contains our `COMMENT`.
-                # Otherwise we need to search it in the raw source code :/
-                if ut_method.has_args():
-                    reason = next(
-                        (
-                            node.value
-                            for node in ast.walk(dec_node)
-                            if isinstance(node, ast.Constant)
-                            and isinstance(node.value, str)
-                            and COMMENT in node.value
-                        ),
-                        None,
-                    )
-
-                    # If we didn't find a constant containing <COMMENT>,
-                    # then we didn't put this decorator
-                    if not reason:
-                        continue
-
-                    if ut_method.has_cond():
-                        cond = ast.unparse(dec_node.args[0])
-                else:
-                    pattern = re.compile(rf"{COMMENT}.?(.*)")
-                    dec_lineno = dec_node.lineno
-
-                    curr_line = lines[dec_lineno - 1]
-                    prev_line = lines[dec_lineno - 2]
-
-                    # If we see our comment at the decorator line, take it
-                    if found := pattern.search(curr_line):
-                        reason = found.group()
-                    elif prev_line.strip().startswith("#") and (
-                        found := pattern.search(prev_line)
-                    ):
-                        # Search the previous line of the decorator,
-                        # only take the comment if the line starts with a `#`
-                        reason = found.group()
-                    else:
-                        # Didn't find our `COMMENT`, so the patch isn't ours :)
-                        continue
-
-                reason = reason.removeprefix(COMMENT).strip(";:, ")
-                spec = PatchSpec(ut_method, cond, reason)
-                yield cls(parent_class, fn_node.name, spec)
+        visitor = PatchEntryVisitor(lines)
+        visitor.visit(tree)
+        yield from visitor.patches
 
 
 def iter_tests(
@@ -251,6 +265,15 @@ def extract_patches(contents: str) -> Patches:
     return build_patch_dict(iter_patches(contents))
 
 
+def modification_from_node_specs(node, specs):
+    lineno = min(
+        (dec_node.lineno for dec_node in node.decorator_list), default=node.lineno
+    )
+    indent = " " * node.col_offset
+    patch_lines = "\n".join(spec.as_decorator() for spec in specs)
+    return (lineno - 1, textwrap.indent(patch_lines, indent))
+
+
 def _iter_patch_lines(
     tree: ast.Module, patches: Patches
 ) -> "Iterator[tuple[int, str]]":
@@ -262,7 +285,15 @@ def _iter_patch_lines(
     async_methods: dict[str, set[str]] = {}
     # Track class bases for inherited async method lookup
     class_bases: dict[str, list[str]] = {}
-    all_classes = {node.name for node in tree.body if isinstance(node, ast.ClassDef)}
+    all_classes = set()
+    all_class_nodes = []
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+
+        all_classes.add(node.name)
+        all_class_nodes.append(node)
+
     for node in tree.body:
         if isinstance(node, ast.ClassDef):
             cache[node.name] = node.end_lineno
@@ -284,13 +315,7 @@ def _iter_patch_lines(
         if not specs:
             continue
 
-        lineno = min(
-            (dec_node.lineno for dec_node in fn_node.decorator_list),
-            default=fn_node.lineno,
-        )
-        indent = " " * fn_node.col_offset
-        patch_lines = "\n".join(spec.as_decorator() for spec in specs)
-        yield (lineno - 1, textwrap.indent(patch_lines, indent))
+        yield modification_from_node_specs(fn_node, specs)
 
     # Phase 2: Iterate and mark inherited tests
     for cls_name, tests in sorted(patches.items()):
@@ -300,6 +325,10 @@ def _iter_patch_lines(
             continue
 
         for test_name, specs in sorted(tests.items()):
+            if test_name == "__self__":
+                # Yielding modifications for the class itself should be done during phase 3
+                continue
+
             decorators = "\n".join(spec.as_decorator() for spec in specs)
             # Check current class and ancestors for async method
             is_async = False
@@ -314,6 +343,7 @@ def _iter_patch_lines(
                     is_async = True
                     break
                 queue.extend(class_bases.get(cur, []))
+
             if is_async:
                 patch_lines = f"""
 {decorators}
@@ -328,6 +358,11 @@ def {test_name}(self):
 """.rstrip()
             yield (lineno, textwrap.indent(patch_lines, DEFAULT_INDENT))
 
+    # Phase 3: Mark the class itself
+    for cls_node in all_class_nodes:
+        if cls_specs := patches.get(cls_node.name, {}).pop("__self__", None):
+            yield modification_from_node_specs(cls_node, cls_specs)
+
 
 def _has_unittest_import(tree: ast.Module) -> bool:
     """Check if 'import unittest' is already present in the file."""
@@ -406,3 +441,42 @@ def patches_from_json(data: dict) -> Patches:
         }
         for cls_name, tests in data.items()
     }
+
+
+def _single_to_double_quotes(s: str) -> str:
+    """
+    Convert single-quoted strings to double-quoted strings.
+
+    Falls back to original if conversion breaks the AST equivalence.
+    """
+    import re
+
+    def replace_string(match: re.Match) -> str:
+        content = match.group(1)
+        # Unescape single quotes and escape double quotes
+        content = content.replace("\\'", "'").replace('"', '\\"')
+        return f'"{content}"'
+
+    # Match single-quoted strings (handles escaped single quotes inside)
+    converted = re.sub(r"'((?:[^'\\]|\\.)*)'", replace_string, s)
+
+    # Verify: parse converted and unparse should equal original
+    try:
+        converted_ast = ast.parse(converted, mode="eval")
+        if ast.unparse(converted_ast) == s:
+            return converted
+    except SyntaxError:
+        pass
+
+    # Fall back to original if conversion failed
+    return s
+
+
+@contextlib.contextmanager
+def temp_attr(obj: object, attr: str, value: object):
+    old = getattr(obj, attr, None)
+    setattr(obj, attr, value)
+    try:
+        yield obj
+    finally:
+        setattr(obj, attr, old)
diff --git a/scripts/update_lib/tests/test_patch_spec.py b/scripts/update_lib/tests/test_patch_spec.py
index 2598d25b3..92a0a35a1 100644
--- a/scripts/update_lib/tests/test_patch_spec.py
+++ b/scripts/update_lib/tests/test_patch_spec.py
@@ -345,6 +345,32 @@ class TestFoo(unittest.TestCase):
         self.assertIn("@unittest.expectedFailure", result)
         self.assertIn(COMMENT, result)
 
+    def test_round_trip_with_patches_on_class(self):
+        """Test that extracted patches can be re-applied."""
+        original = f"""import unittest
+
+@unittest.skipIf(a == b, "{COMMENT}")
+@unittest.expectedFailure  # {COMMENT}
+class TestFoo(unittest.TestCase):
+    ...
+"""
+        # Extract patches
+        patches = extract_patches(original)
+
+        # Apply to clean code
+        clean = """import unittest
+
+class TestFoo(unittest.TestCase):
+    def test_one(self):
+        pass
+"""
+        result = apply_patches(clean, patches)
+
+        # Should have the decorator
+        self.assertIn("@unittest.expectedFailure", result)
+        self.assertIn("@unittest.skipIf", result)
+        self.assertIn(COMMENT, result)
+
 
 class TestFindImportInsertLine(unittest.TestCase):
     """Tests for _find_import_insert_line function."""