#!/usr/bin/env python

"""

Quick update for test files from CPython.

Usage:

# Library + test: copy lib, then patch + auto-mark test + commit

python scripts/update_lib quick cpython/Lib/dataclasses.py

# Shortcut: just the module name

python scripts/update_lib quick dataclasses

# Test file: patch + auto-mark

python scripts/update_lib quick cpython/Lib/test/test_foo.py

# Test file: migrate only

python scripts/update_lib quick cpython/Lib/test/test_foo.py --no-auto-mark

# Test file: auto-mark only (Lib/ path implies --no-migrate)

python scripts/update_lib quick Lib/test/test_foo.py

# Directory: patch all + auto-mark all

python scripts/update_lib quick cpython/Lib/test/test_dataclasses/

# Skip git commit

python scripts/update_lib quick dataclasses --no-commit

"""

import argparse

import pathlib

import sys

sys.path.insert(0, str(pathlib.Path(__file__).parent.parent))

from update_lib.deps import DEPENDENCIES, get_test_paths

from update_lib.file_utils import (

construct_lib_path,

get_cpython_dir,

get_module_name,

get_test_files,

is_lib_path,

is_test_path,

lib_to_test_path,

parse_lib_path,

resolve_module_path,

safe_read_text,

)

def collect_original_methods(

lib_path: pathlib.Path,

) -> set[tuple[str, str]] | dict[pathlib.Path, set[tuple[str, str]]] | None:

"""

Collect original test methods from lib path before patching.

Returns:

- For file: set of (class_name, method_name) or None if file doesn't exist

- For directory: dict mapping file path to set of methods, or None if dir doesn't exist

"""

from update_lib.cmd_auto_mark import extract_test_methods

if not lib_path.exists():

return None

if lib_path.is_file():

content = safe_read_text(lib_path)

return extract_test_methods(content) if content else set()

else:

result = {}

for lib_file in get_test_files(lib_path):

content = safe_read_text(lib_file)

if content:

result[lib_file.resolve()] = extract_test_methods(content)

return result

def quick(

src_path: pathlib.Path,

no_migrate: bool = False,

no_auto_mark: bool = False,

mark_failure: bool = False,

verbose: bool = True,

skip_build: bool = False,

) -> list[pathlib.Path]:

"""

Process a file or directory: migrate + auto-mark.

Args:

src_path: Source path (file or directory)

no_migrate: Skip migration step

no_auto_mark: Skip auto-mark step

mark_failure: Add @expectedFailure to ALL failing tests

verbose: Print progress messages

skip_build: Skip cargo build, use pre-built binary

Returns:

List of extra paths (data dirs, hard deps) that were copied/migrated.

"""

from update_lib.cmd_auto_mark import auto_mark_directory, auto_mark_file

from update_lib.cmd_migrate import patch_directory, patch_file

extra_paths: list[pathlib.Path] = []

# Determine lib_path and whether to migrate

if is_lib_path(src_path):

no_migrate = True

lib_path = src_path

else:

lib_path = parse_lib_path(src_path)

is_dir = src_path.is_dir()

# Capture original test methods before migration (for smart auto-mark)

original_methods = collect_original_methods(lib_path)

# Step 1: Migrate

if not no_migrate:

if is_dir:

patch_directory(src_path, lib_path, verbose=verbose)

else:

patch_file(src_path, lib_path, verbose=verbose)

# Step 1.5: Handle test dependencies

from update_lib.deps import get_test_dependencies

test_deps = get_test_dependencies(src_path)

# Migrate dependency files

for dep_src in test_deps["hard_deps"]:

dep_lib = parse_lib_path(dep_src)

if verbose:

print(f"Migrating dependency: {dep_src.name}")

if dep_src.is_dir():

patch_directory(dep_src, dep_lib, verbose=False)

else:

patch_file(dep_src, dep_lib, verbose=False)

extra_paths.append(dep_lib)

# Copy data directories (no migration)

import shutil

for data_src in test_deps["data"]:

data_lib = parse_lib_path(data_src)

if verbose:

print(f"Copying data: {data_src.name}")

if data_lib.exists():

if data_lib.is_dir():

shutil.rmtree(data_lib)

else:

data_lib.unlink()

if data_src.is_dir():

shutil.copytree(data_src, data_lib)

else:

data_lib.parent.mkdir(parents=True, exist_ok=True)

shutil.copy2(data_src, data_lib)

extra_paths.append(data_lib)

# Step 2: Auto-mark

if not no_auto_mark:

if not lib_path.exists():

raise FileNotFoundError(f"Path not found: {lib_path}")

if is_dir:

num_added, num_removed, _ = auto_mark_directory(

lib_path,

mark_failure=mark_failure,

verbose=verbose,

original_methods_per_file=original_methods,

skip_build=skip_build,

)

else:

num_added, num_removed, _ = auto_mark_file(

lib_path,

mark_failure=mark_failure,

verbose=verbose,

original_methods=original_methods,

skip_build=skip_build,

)

if verbose:

if num_added:

print(f"Added expectedFailure to {num_added} tests")

print(f"Removed expectedFailure from {num_removed} tests")

return extra_paths

def get_cpython_version(cpython_dir: pathlib.Path) -> str:

"""Get CPython version from git tag."""

import subprocess

result = subprocess.run(

["git", "describe", "--tags"],

cwd=cpython_dir,

capture_output=True,

text=True,

check=True,

)

return result.stdout.strip()

def git_commit(

name: str,

lib_path: pathlib.Path | None,

test_paths: list[pathlib.Path] | pathlib.Path | None,

cpython_dir: pathlib.Path,

hard_deps: list[pathlib.Path] | None = None,

verbose: bool = True,

) -> bool:

"""Commit changes with CPython author.

Args:

name: Module name (e.g., "dataclasses")

lib_path: Path to library file/directory (or None)

test_paths: Path(s) to test file/directory (or None)

cpython_dir: Path to cpython directory

hard_deps: Path(s) to hard dependency files (or None)

verbose: Print progress messages

Returns:

True if commit was created, False otherwise

"""

import subprocess

# Normalize test_paths to list

if test_paths is None:

test_paths = []

elif isinstance(test_paths, pathlib.Path):

test_paths = [test_paths]

# Normalize hard_deps to list

if hard_deps is None:

hard_deps = []

# Stage changes

paths_to_add = []

if lib_path and lib_path.exists():

paths_to_add.append(str(lib_path))

for test_path in test_paths:

if test_path and test_path.exists():

paths_to_add.append(str(test_path))

for dep_path in hard_deps:

if dep_path and dep_path.exists():

paths_to_add.append(str(dep_path))

if not paths_to_add:

return False

version = get_cpython_version(cpython_dir)

subprocess.run(["git", "add"] + paths_to_add, check=True)

# Check if there are staged changes

result = subprocess.run(

["git", "diff", "--cached", "--quiet"],

capture_output=True,

)

if result.returncode == 0:

if verbose:

print("No changes to commit")

return False

# Commit with CPython author

message = f"Update {name} from {version}"

subprocess.run(

[

"git",

"commit",

"--author",

"CPython Developers <>",

"-m",

message,

],

check=True,

)

if verbose:

print(f"Committed: {message}")

return True

def _expand_shortcut(path: pathlib.Path) -> pathlib.Path:

"""Expand simple name to cpython/Lib path if it exists.

Examples:

dataclasses -> cpython/Lib/dataclasses.py (if exists)

json -> cpython/Lib/json/ (if exists)

test_types -> cpython/Lib/test/test_types.py (if exists)

regrtest -> cpython/Lib/test/libregrtest (from DEPENDENCIES)

"""

# Only expand if it's a simple name (no path separators) and doesn't exist

if "/" in str(path) or path.exists():

return path

name = str(path)

# Check DEPENDENCIES table for path overrides (e.g., regrtest)

from update_lib.deps import DEPENDENCIES

if name in DEPENDENCIES and "lib" in DEPENDENCIES[name]:

lib_paths = DEPENDENCIES[name]["lib"]

if lib_paths:

override_path = construct_lib_path("cpython", lib_paths[0])

if override_path.exists():

return override_path

# Test shortcut: test_foo -> cpython/Lib/test/test_foo

if name.startswith("test_"):

resolved = resolve_module_path(f"test/{name}", "cpython", prefer="dir")

if resolved.exists():

return resolved

# Library shortcut: foo -> cpython/Lib/foo

resolved = resolve_module_path(name, "cpython", prefer="file")

if resolved.exists():

return resolved

# Extension module shortcut: winreg -> cpython/Lib/test/test_winreg

# For C/Rust extension modules that have no Python source but have tests

resolved = resolve_module_path(f"test/test_{name}", "cpython", prefer="dir")

if resolved.exists():

return resolved

# Return original (will likely fail later with a clear error)

return path

def main(argv: list[str] | None = None) -> int:

parser = argparse.ArgumentParser(

description=__doc__,

formatter_class=argparse.RawDescriptionHelpFormatter,

)

parser.add_argument(

"path",

type=pathlib.Path,

help="Source path (file or directory)",

)

parser.add_argument(

"--copy",

action=argparse.BooleanOptionalAction,

default=True,

help="Copy library file (default: enabled, implied disabled if test path)",

)

parser.add_argument(

"--migrate",

action=argparse.BooleanOptionalAction,

default=True,

help="Migrate test file (default: enabled, implied disabled if Lib/ path)",

)

parser.add_argument(

"--auto-mark",

action=argparse.BooleanOptionalAction,

default=True,

help="Auto-mark test failures (default: enabled)",

)

parser.add_argument(

"--mark-failure",

action="store_true",

help="Add @expectedFailure to failing tests",

)

parser.add_argument(

"--commit",

action=argparse.BooleanOptionalAction,

default=True,

help="Create git commit (default: enabled)",

)

parser.add_argument(

"--build",

action=argparse.BooleanOptionalAction,

default=True,

help="Build with cargo (default: enabled)",

)

args = parser.parse_args(argv)

try:

src_path = args.path

# Shortcut: expand simple name to cpython/Lib path

src_path = _expand_shortcut(src_path)

original_src = src_path # Keep for commit

# Track library path for commit

lib_file_path = None

test_path = None

hard_deps_for_commit = []

# If it's a library path (not test path), do copy_lib first

if not is_test_path(src_path):

# Get library destination path for commit

lib_file_path = parse_lib_path(src_path)

if args.copy:

from update_lib.cmd_copy_lib import copy_lib

copy_lib(src_path)

# Get all test paths from DEPENDENCIES (or fall back to default)

module_name = get_module_name(original_src)

cpython_dir = get_cpython_dir(original_src)

test_src_paths = get_test_paths(module_name, str(cpython_dir))

# Fall back to default test path if DEPENDENCIES has no entry

if not test_src_paths:

default_test = lib_to_test_path(original_src)

if default_test.exists():

test_src_paths = (default_test,)

# Collect hard dependencies for commit

lib_deps = DEPENDENCIES.get(module_name, {})

for dep_name in lib_deps.get("hard_deps", []):

dep_lib_path = pathlib.Path("Lib") / dep_name

if dep_lib_path.exists():

hard_deps_for_commit.append(dep_lib_path)

# Process all test paths

test_paths_for_commit = []

for test_src in test_src_paths:

if not test_src.exists():

print(f"Warning: Test path does not exist: {test_src}")

continue

test_lib_path = parse_lib_path(test_src)

test_paths_for_commit.append(test_lib_path)

extra = quick(

test_src,

no_migrate=not args.migrate,

no_auto_mark=not args.auto_mark,

mark_failure=args.mark_failure,

skip_build=not args.build,

)

hard_deps_for_commit.extend(extra)

test_paths = test_paths_for_commit

else:

# It's a test path - process single test

test_path = (

parse_lib_path(src_path) if not is_lib_path(src_path) else src_path

)

extra = quick(

src_path,

no_migrate=not args.migrate,

no_auto_mark=not args.auto_mark,

mark_failure=args.mark_failure,

skip_build=not args.build,

)

hard_deps_for_commit.extend(extra)

test_paths = [test_path]

# Step 3: Git commit

if args.commit:

cpython_dir = get_cpython_dir(original_src)

git_commit(

get_module_name(original_src),

lib_file_path,

test_paths,

cpython_dir,

hard_deps=hard_deps_for_commit,

)

return 0

except ValueError as e:

print(f"Error: {e}", file=sys.stderr)

return 1

except FileNotFoundError as e:

print(f"Error: {e}", file=sys.stderr)

return 1

except Exception as e:

# Handle TestRunError with a clean message

from update_lib.cmd_auto_mark import TestRunError

if isinstance(e, TestRunError):

print(f"Error: {e}", file=sys.stderr)

return 1

raise

if __name__ == "__main__":

sys.exit(main())