Support scoped canonical test updates in update_lib

Index conditional test classes and methods without crossing namespaces, reject ambiguous targets, and preserve inherited override indentation. Add hash-pinned whole-test copying that verifies execution context and exact source bytes, with focused Python fixtures.

Assisted-by: OpenAI Codex:model identifier unavailable

scripts/update_lib/__main__.py

@@ -8,6 +8,7 @@ Usage:
     python scripts/update_lib migrate cpython/Lib/test/test_foo.py
     python scripts/update_lib patches --from Lib/test/foo.py --to cpython/Lib/test/foo.py
     python scripts/update_lib auto-mark Lib/test/test_foo.py
+    python scripts/update_lib copy-test SOURCE DESTINATION --test ClassName.method_name --source-sha256 HASH --destination-sha256 HASH
 """
 
 import argparse
@@ -49,6 +50,11 @@ def main(argv: list[str] | None = None) -> int:
         help="Copy library file/directory from CPython (delete existing first)",
         add_help=False,
     )
+    subparsers.add_parser(
+        "copy-test",
+        help="Copy one complete canonical test from hash-pinned inputs",
+        add_help=False,
+    )
     subparsers.add_parser(
         "deps",
         help="Show dependency information for a module",
@@ -72,6 +78,11 @@ def main(argv: list[str] | None = None) -> int:
 
         return copy_lib_main(remaining)
 
+    if args.command == "copy-test":
+        from update_lib.cmd_copy_test import main as copy_test_main
+
+        return copy_test_main(remaining)
+
     if args.command == "migrate":
         from update_lib.cmd_migrate import main as migrate_main
 

scripts/update_lib/cmd_copy_test.py

@@ -0,0 +1,125 @@
+"""Copy one complete canonical test definition without refreshing its module."""
+
+import argparse
+import ast
+import hashlib
+import json
+from pathlib import Path
+
+from update_lib.patch_spec import _iter_class_methods, _iter_module_classes
+
+
+def _definition(contents: str, qualified_name: str):
+    parts = qualified_name.split(".")
+    if len(parts) != 2 or not all(part.isidentifier() for part in parts):
+        raise ValueError("test name must be ClassName.method_name")
+    tree = ast.parse(contents)
+    classes = [node for node in _iter_module_classes(tree) if node.name == parts[0]]
+    if len(classes) != 1:
+        raise ValueError(f"Expected one class named {parts[0]}, found {len(classes)}")
+    methods = [
+        node for node in _iter_class_methods(classes[0]) if node.name == parts[1]
+    ]
+    if len(methods) != 1:
+        raise ValueError(
+            f"Expected one method named {qualified_name}, found {len(methods)}"
+        )
+    method = methods[0]
+    parents = {
+        child: parent
+        for parent in ast.walk(tree)
+        for child in ast.iter_child_nodes(parent)
+    }
+    guards = []
+    child = method
+    while child in parents:
+        parent = parents[child]
+        if isinstance(parent, ast.If):
+            branch = "body" if child in parent.body else "orelse"
+            guards.append((ast.dump(parent.test), branch))
+        elif not isinstance(parent, (ast.Module, ast.ClassDef)):
+            raise ValueError("Only module/class scopes and if guards are supported")
+        child = parent
+    start = min([method.lineno, *(node.lineno for node in method.decorator_list)]) - 1
+    return method, guards, start, method.end_lineno
+
+
+def copy_test(source: str, destination: str, qualified_name: str) -> str:
+    """Replace an unambiguous whole definition, preserving all surrounding bytes.
+
+    Guard, signature, decorator and indentation differences require a broader
+    migration rather than silently changing the selected test's execution context.
+    """
+    src, src_guards, src_start, src_end = _definition(source, qualified_name)
+    dst, dst_guards, dst_start, dst_end = _definition(destination, qualified_name)
+    if src_guards != dst_guards:
+        raise ValueError("Source and destination test guards differ")
+
+    def signature(node):
+        return (
+            ast.dump(node.args),
+            ast.dump(node.returns) if node.returns is not None else None,
+            [ast.dump(param) for param in getattr(node, "type_params", [])],
+        )
+
+    if type(src) is not type(dst) or signature(src) != signature(dst):
+        raise ValueError("Source and destination test signatures differ")
+    if [ast.dump(node) for node in src.decorator_list] != [
+        ast.dump(node) for node in dst.decorator_list
+    ]:
+        raise ValueError(
+            "Source and destination decorators differ; preserve markers explicitly"
+        )
+    if src.col_offset != dst.col_offset:
+        raise ValueError("Source and destination indentation differs")
+    source_lines = source.splitlines(keepends=True)
+    destination_lines = destination.splitlines(keepends=True)
+    definition = source_lines[src_start:src_end]
+    result = "".join(
+        destination_lines[:dst_start] + definition + destination_lines[dst_end:]
+    )
+    _, _, result_start, result_end = _definition(result, qualified_name)
+    if result.splitlines(keepends=True)[result_start:result_end] != definition:
+        raise ValueError("Copied definition is not byte-identical to the source")
+    return result
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("source", type=Path)
+    parser.add_argument("destination", type=Path)
+    parser.add_argument("--test", required=True, help="Exact ClassName.method_name")
+    parser.add_argument("--source-sha256", required=True)
+    parser.add_argument("--destination-sha256", required=True)
+    args = parser.parse_args(argv)
+    try:
+        source = args.source.read_bytes()
+        destination = args.destination.read_bytes()
+        for label, contents, expected in (
+            ("source", source, args.source_sha256),
+            ("destination", destination, args.destination_sha256),
+        ):
+            if hashlib.sha256(contents).hexdigest() != expected:
+                raise ValueError(f"{label} SHA256 does not match the pinned input")
+        result = copy_test(
+            source.decode("utf-8"), destination.decode("utf-8"), args.test
+        ).encode("utf-8")
+        # Recheck the destination immediately before writing so stale invocations
+        # fail instead of overwriting an intervening edit.
+        if args.destination.read_bytes() != destination:
+            raise ValueError("Destination changed during test preparation")
+        args.destination.write_bytes(result)
+    except (OSError, UnicodeError, SyntaxError, ValueError) as exc:
+        parser.exit(1, f"Error: {exc}\n")
+    print(
+        json.dumps(
+            {
+                "test": args.test,
+                "source_sha256": args.source_sha256,
+                "destination_before_sha256": args.destination_sha256,
+                "destination_after_sha256": hashlib.sha256(result).hexdigest(),
+            },
+            sort_keys=True,
+        )
+    )
+    return 0

scripts/update_lib/patch_spec.py

@@ -197,7 +197,10 @@ class PatchEntryVisitor(ast.NodeVisitor):
                 patch = patch._replace(test_name="__self__")
                 self.patches.append(patch)
 
-            self.generic_visit(node)
+            # Keep conditional methods in this namespace, but do not enter
+            # nested classes or the bodies of methods.
+            for item in _iter_class_methods(node):
+                self.visit(item)
 
 
 class PatchEntry(typing.NamedTuple):
@@ -228,22 +231,35 @@ class PatchEntry(typing.NamedTuple):
         yield from visitor.patches
 
 
+def _iter_module_classes(node: ast.AST) -> "Iterator[ast.ClassDef]":
+    """Find module-scope classes, including conditional definitions.
+
+    Function-local and class-nested definitions have different namespaces and
+    cannot be addressed by the patch format's unqualified class names.
+    """
+    for child in ast.iter_child_nodes(node):
+        if isinstance(child, ast.ClassDef):
+            yield child
+        elif not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
+            yield from _iter_module_classes(child)
+
+
+def _iter_class_methods(
+    node: ast.AST,
+) -> "Iterator[ast.FunctionDef | ast.AsyncFunctionDef]":
+    for child in ast.iter_child_nodes(node):
+        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            yield child
+        elif not isinstance(child, (ast.ClassDef, ast.Lambda)):
+            yield from _iter_class_methods(child)
+
+
 def iter_tests(
     tree: ast.Module,
 ) -> "Iterator[tuple[ast.ClassDef, ast.FunctionDef | ast.AsyncFunctionDef]]":
-    for key, nodes in ast.iter_fields(tree):
-        if key != "body":
-            continue
-
-        for cls_node in nodes:
-            if not isinstance(cls_node, ast.ClassDef):
-                continue
-
-            for fn_node in cls_node.body:
-                if not isinstance(fn_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                    continue
-
-                yield (cls_node, fn_node)
+    for cls_node in _iter_module_classes(tree):
+        for fn_node in _iter_class_methods(cls_node):
+            yield (cls_node, fn_node)
 
 
 def iter_patches(contents: str) -> "Iterator[PatchEntry]":
@@ -286,24 +302,28 @@ def _iter_patch_lines(
     # Track class bases for inherited async method lookup
     class_bases: dict[str, list[str]] = {}
     all_classes = set()
-    all_class_nodes = []
-    for node in tree.body:
-        if not isinstance(node, ast.ClassDef):
-            continue
-
+    all_class_nodes = list(_iter_module_classes(tree))
+    for node in all_class_nodes:
+        if node.name in all_classes:
+            raise ValueError(f"Ambiguous module-scope class name: {node.name}")
         all_classes.add(node.name)
-        all_class_nodes.append(node)
+        method_counts = collections.Counter(
+            item.name for item in _iter_class_methods(node)
+        )
+        for method in patches.get(node.name, {}):
+            if method_counts[method] > 1:
+                raise ValueError(f"Ambiguous method name: {node.name}.{method}")
 
-    for node in tree.body:
+    for node in all_class_nodes:
         if isinstance(node, ast.ClassDef):
-            cache[node.name] = node.end_lineno
+            cache[node.name] = (node.end_lineno, node.col_offset + len(DEFAULT_INDENT))
             class_bases[node.name] = [
                 base.id
                 for base in node.bases
                 if isinstance(base, ast.Name) and base.id in all_classes
             ]
             cls_async: set[str] = set()
-            for item in node.body:
+            for item in _iter_class_methods(node):
                 if isinstance(item, ast.AsyncFunctionDef):
                     cls_async.add(item.name)
             if cls_async:
@@ -319,10 +339,11 @@ def _iter_patch_lines(
 
     # Phase 2: Iterate and mark inherited tests
     for cls_name, tests in sorted(patches.items()):
-        lineno = cache.get(cls_name)
-        if not lineno:
+        location = cache.get(cls_name)
+        if not location:
             print(f"WARNING: {cls_name} does not exist in remote file", file=sys.stderr)
             continue
+        lineno, indent = location
 
         for test_name, specs in sorted(tests.items()):
             if test_name == "__self__":
@@ -356,7 +377,7 @@ async def {test_name}(self):
 def {test_name}(self):
 {DEFAULT_INDENT}return super().{test_name}()
 """.rstrip()
-            yield (lineno, textwrap.indent(patch_lines, DEFAULT_INDENT))
+            yield (lineno, textwrap.indent(patch_lines, " " * indent))
 
     # Phase 3: Mark the class itself
     for cls_node in all_class_nodes:

scripts/update_lib/tests/test_copy_test.py

@@ -0,0 +1,121 @@
+"""Tests for bounded, hash-pinned canonical test copying."""
+
+import ast
+import contextlib
+import hashlib
+import io
+import tempfile
+import unittest
+from pathlib import Path
+
+from update_lib.cmd_copy_test import copy_test, main
+
+
+class TestCopyTest(unittest.TestCase):
+    destination = '''# Destination-only header
+class Finder:
+    if int.__doc__:
+        def example():
+            r"""Count is below 800. π\\n"""
+
+    def untouched(self):
+        return "destination"
+'''
+    source = '''# Canonical header must not be copied
+class Finder:
+    if int.__doc__:
+        def example():
+            r"""Count is below 850. π\\n"""
+
+    def untouched(self):
+        return "canonical"
+'''
+
+    def test_complete_definition_and_surroundings(self):
+        result = copy_test(self.source, self.destination, "Finder.example")
+        self.assertEqual(result, self.destination.replace("below 800", "below 850"))
+        self.assertIn('return "destination"', result)
+        source_method = ast.parse(self.source).body[0].body[0].body[0]
+        result_method = ast.parse(result).body[0].body[0].body[0]
+        self.assertEqual(
+            ast.get_source_segment(self.source, source_method),
+            ast.get_source_segment(result, result_method),
+        )
+
+    def test_guard_mismatch_is_rejected(self):
+        with self.assertRaisesRegex(ValueError, "guards differ"):
+            copy_test(
+                self.source.replace("if int.__doc__", "if enabled"),
+                self.destination,
+                "Finder.example",
+            )
+
+    def test_missing_or_ambiguous_definition_is_rejected(self):
+        for source in ("class Other: pass\n", self.source + self.source):
+            with self.subTest(source=source):
+                with self.assertRaisesRegex(ValueError, "Expected one class"):
+                    copy_test(source, self.destination, "Finder.example")
+        duplicate = self.source.replace(
+            "    def untouched(self):",
+            "    def example():\n        pass\n\n    def untouched(self):",
+        )
+        with self.assertRaisesRegex(ValueError, "Expected one method"):
+            copy_test(duplicate, self.destination, "Finder.example")
+
+    def test_decorator_and_signature_changes_are_rejected(self):
+        decorated = self.source.replace(
+            "        def example():", "        @decorator\n        def example():"
+        )
+        with self.assertRaisesRegex(ValueError, "decorators differ"):
+            copy_test(decorated, self.destination, "Finder.example")
+        with self.assertRaisesRegex(ValueError, "signatures differ"):
+            copy_test(
+                self.source.replace("def example():", "def example(arg):"),
+                self.destination,
+                "Finder.example",
+            )
+
+    def test_nested_local_scope_is_not_selected(self):
+        local = "def factory():\n" + "".join(
+            "    " + line for line in self.source.splitlines(True)
+        )
+        with self.assertRaisesRegex(ValueError, "Expected one class"):
+            copy_test(local, self.destination, "Finder.example")
+
+    def test_return_annotation_and_type_parameters_are_rejected(self):
+        for declaration in ("def example() -> int:", "def example[T]():"):
+            with self.subTest(declaration=declaration):
+                source = self.source.replace("def example():", declaration)
+                with self.assertRaisesRegex(ValueError, "signatures differ"):
+                    copy_test(source, self.destination, "Finder.example")
+
+    def test_cli_hashes_and_no_write_on_stale_input(self):
+        with tempfile.TemporaryDirectory() as temp:
+            source = Path(temp) / "source.py"
+            destination = Path(temp) / "destination.py"
+            source.write_bytes(self.source.encode())
+            destination.write_bytes(self.destination.encode())
+            args = [
+                str(source),
+                str(destination),
+                "--test",
+                "Finder.example",
+                "--source-sha256",
+                hashlib.sha256(source.read_bytes()).hexdigest(),
+                "--destination-sha256",
+                "0" * 64,
+            ]
+            with (
+                contextlib.redirect_stderr(io.StringIO()),
+                self.assertRaises(SystemExit),
+            ):
+                main(args)
+            self.assertEqual(destination.read_bytes(), self.destination.encode())
+            args[-1] = hashlib.sha256(destination.read_bytes()).hexdigest()
+            with contextlib.redirect_stdout(io.StringIO()) as output:
+                self.assertEqual(main(args), 0)
+            self.assertIn('"source_sha256"', output.getvalue())
+            self.assertEqual(
+                destination.read_text(),
+                self.destination.replace("below 800", "below 850"),
+            )

scripts/update_lib/tests/test_patch_spec.py

@@ -67,6 +67,38 @@ class TestAsync(unittest.TestCase):
 class TestExtractPatches(unittest.TestCase):
     """Tests for extract_patches function."""
 
+    def test_class_conditional_method_markers_are_preserved(self):
+        code = """class Owner:
+    if enabled:
+        @unittest.expectedFailure  # TODO: RUSTPYTHON
+        def test_value(self):
+            return 42
+"""
+        patches = extract_patches(code)
+        self.assertIn("test_value", patches["Owner"])
+        clean = code.replace(
+            "        @unittest.expectedFailure  # TODO: RUSTPYTHON\n", ""
+        )
+        result = apply_patches(clean, patches)
+        self.assertIn("        @unittest.expectedFailure", result)
+        method = list(iter_tests(ast.parse(result)))[0][1]
+        self.assertEqual(method.col_offset, 8)
+        self.assertEqual(ast.dump(method.body[0]), "Return(value=Constant(value=42))")
+
+    def test_nested_class_markers_do_not_target_module_names(self):
+        code = """class Container:
+    class Target:
+        @unittest.expectedFailure  # TODO: RUSTPYTHON
+        def test_value(self):
+            pass
+def factory():
+    class Target:
+        @unittest.expectedFailure  # TODO: RUSTPYTHON
+        def test_value(self):
+            pass
+"""
+        self.assertEqual(extract_patches(code), {})
+
     def test_extract_expected_failure(self):
         """Test extracting @unittest.expectedFailure decorator."""
         code = f"""
@@ -177,6 +209,106 @@ class TestFoo(unittest.TestCase):
 class TestApplyPatches(unittest.TestCase):
     """Tests for apply_patches function."""
 
+    def test_conditional_class_existing_method(self):
+        code = """import unittest
+if enabled:
+    class TestConditional:
+        def test_value(self):
+            assert 1 + 1 == 2
+"""
+        result = apply_patches(
+            code,
+            {
+                "TestConditional": {
+                    "test_value": [
+                        PatchSpec(
+                            UtMethod.ExpectedFailureIf, "condition", "known limitation"
+                        )
+                    ]
+                }
+            },
+        )
+        original = list(iter_tests(ast.parse(code)))[0][1]
+        updated = list(iter_tests(ast.parse(result)))[0][1]
+        self.assertEqual(ast.dump(original.body[0]), ast.dump(updated.body[0]))
+        self.assertEqual(updated.col_offset, 8)
+        self.assertIn("        @unittest.expectedFailureIf", result)
+
+    def test_conditional_class_inherited_override(self):
+        code = """import unittest
+class Parent:
+    async def test_value(self):
+        return 42
+if enabled:
+    class Child(Parent):
+        pass
+"""
+        result = apply_patches(
+            code, {"Child": {"test_value": [PatchSpec(UtMethod.ExpectedSuccess)]}}
+        )
+        methods = list(iter_tests(ast.parse(result)))
+        self.assertEqual([cls.name for cls, _ in methods], ["Parent", "Child"])
+        self.assertIsInstance(methods[1][1], ast.AsyncFunctionDef)
+        self.assertIn("        return await super().test_value()", result)
+        self.assertIn("        @unittest.expectedSuccess", result)
+
+    def test_duplicate_conditional_class_names_are_rejected(self):
+        code = """if enabled:
+    class Duplicate:
+        pass
+else:
+    class Duplicate:
+        pass
+"""
+        with self.assertRaisesRegex(ValueError, "Ambiguous.*Duplicate"):
+            apply_patches(
+                code,
+                {"Duplicate": {"test_value": [PatchSpec(UtMethod.ExpectedFailure)]}},
+            )
+
+    def test_local_class_names_do_not_shadow_module_class(self):
+        code = """import unittest
+class Target:
+    def test_value(self):
+        return 1
+class Container:
+    class Target:
+        def test_value(self):
+            return 2
+def factory():
+    class Target:
+        def test_value(self):
+            return 3
+"""
+        result = apply_patches(
+            code, {"Target": {"test_value": [PatchSpec(UtMethod.ExpectedFailure)]}}
+        )
+        tree = ast.parse(result)
+        targets = [
+            node
+            for node in ast.walk(tree)
+            if isinstance(node, ast.ClassDef) and node.name == "Target"
+        ]
+        self.assertEqual(
+            [len(node.body[0].decorator_list) for node in targets], [1, 0, 0]
+        )
+
+    def test_duplicate_conditional_method_names_are_rejected(self):
+        code = """class Owner:
+    if enabled:
+        def test_value(self):
+            return 1
+    else:
+        def test_value(self):
+            return 2
+"""
+        with self.assertRaisesRegex(
+            ValueError, "Ambiguous method name: Owner.test_value"
+        ):
+            apply_patches(
+                code, {"Owner": {"test_value": [PatchSpec(UtMethod.ExpectedFailure)]}}
+            )
+
     def test_apply_expected_failure(self):
         """Test applying @unittest.expectedFailure."""
         code = """import unittest