From dce6c01583d7fd71c775e6ac6cd3e1af85a4b224 Mon Sep 17 00:00:00 2001 From: Someone Serge Date: Fri, 3 Nov 2023 00:22:34 +0000 Subject: [PATCH] python-prefix-modules: init --- nixpkgs_pytools/prefix_modules.py | 302 ++++++++++++++++++++++++++++++ setup.py | 3 +- tests/test_prefix_modules.py | 118 ++++++++++++ 3 files changed, 422 insertions(+), 1 deletion(-) create mode 100644 nixpkgs_pytools/prefix_modules.py create mode 100644 tests/test_prefix_modules.py diff --git a/nixpkgs_pytools/prefix_modules.py b/nixpkgs_pytools/prefix_modules.py new file mode 100644 index 0000000..a3db08f --- /dev/null +++ b/nixpkgs_pytools/prefix_modules.py @@ -0,0 +1,302 @@ +import os + +if "PREFIX_PYTHON_MODULES_DEBUG" in os.environ: + import pdb + + pdb.set_trace() + +import argparse +import fnmatch +import shutil +import sys +import tempfile +import textwrap +import traceback +from pathlib import Path +from typing import List, Literal, Optional, Tuple + +from rope.base.project import Project +from rope.refactor.move import MoveModule +from rope.refactor.rename import Rename + +parser = argparse.ArgumentParser("prefix-python-modules") +parser.add_argument("repo_root", type=Path) +parser.add_argument( + "--rename-external", + nargs=3, + action="append", + help="Takes (old_name, new_name, glob), e.g.:" + ' --rename-external utils dino.utils "**/extract_dino_features.py"', +) +parser.add_argument("-e", "--exclude-glob", action="append") +parser.add_argument("--dont-catch", action="store_true") +parser.add_argument("--prefix", required=True) +parser.add_argument("--verbose", action="store_true") +parser.add_argument("--quiet", action="store_true") +parser.add_argument( + "--mode", + default="first-error", + choices=("first-error", "keep-going", "interactive"), +) + +CATCH_ERRORS = True + + +def indent(s: str, n: int = 4) -> str: + return textwrap.indent(s, " " * n) + + +def convert_to_packages(project_root, exclude_globs: List[str]): + project = Project(project_root) + python_files = [p for p in project.get_python_files()] + python_files = [ + p for p in python_files if not fnmatch_any_glob(p.path, exclude_globs) + ] + + try: + for f in python_files: + rel_path = Path(f.path) + if rel_path.parent == Path("."): + continue + path = project_root / rel_path.parent / "__init__.py" + path.touch() + finally: + project.close() + + +def apply_changes( + project, + changes, + mode: Literal["first-error", "keep-going", "interactive"], +) -> Tuple[ + Optional[Literal["quit", "next"]], + Optional[str], + Optional[Exception], + Optional[str], +]: + description = changes.get_description() + + while True: + try: + if mode == "interactive": + print("Apply the following patch?") + print(indent(description)) + print("[Y]es, [n]o, [q]uit? [Ynq]") + action = input().lower().strip() + if action == "": + action = "y" + assert action in "ynq", action + if action == "n": + return ("next", None, None, None) + elif action == "q": + return ("quit", None, None, None) + project.do(changes) + project.validate() + except Exception as e: + if not CATCH_ERRORS: + raise + if mode != "interactive": + return (None, description, e, traceback.format_exc()) + + keep_asking = True + action = "q" + while keep_asking: + print( + f"Failed to apply the patch: {e}\n" + "...[r]etry, print [v]erbose error, skip and proceed to the [n]ext patch, or [Q]uit? [rvnQ]" + ) + action = input().lower().strip() + if action == "": + action = "q" + assert action in "rvnq", action + if action == "v": + print(traceback.format_exc()) + continue + keep_asking = False + if action == "q": + return ("quit", description, e, traceback.format_exc()) + elif action == "n": + return ("next", description, e, traceback.format_exc()) + else: + continue + else: + return (None, description, None, None) + + +def fnmatch_any_glob(path, globs): + return any(fnmatch.fnmatchcase(path, g) for g in globs) + + +def prefix_modules( + project_root, + prefix, + *, + mode: Literal["first-error", "keep-going", "interactive"], + verbose: bool, + exclude_globs: List[str], +): + project = Project(project_root) + + parallel_tree_for_rope = tempfile.mkdtemp() + + try: + project.validate() + + python_files = [p for p in project.get_python_files()] + python_files = [p for p in python_files if Path(p.path).parts[0] != prefix] + python_files = [ + p for p in python_files if not fnmatch_any_glob(p.path, exclude_globs) + ] + + toplevel_files = sorted(set(Path(p.path).parts[0] for p in python_files)) + toplevel_module_names = [name.removesuffix(".py") for name in toplevel_files] + + new_package = project.get_folder(prefix) + if not new_package.exists(): + new_package.create() + new_package.create_file("__init__.py") + + successes = [] + failures = [] + for name in toplevel_module_names: + m = project.get_module(name) + r = m.get_resource() + + old_path = r.pathlib + + try: + changes = MoveModule(project, r).get_changes(new_package) + except: + print( + "Fatal error: coudln't prepare a MoveModule for" + f" {name} -> {new_package.path}.{name}" + ) + raise + (action, description, error, tb) = apply_changes( + project, changes, mode=mode + ) + + if error is not None and verbose: + failures.append((description, tb)) + elif error is not None: + failures.append((description, error)) + if error is None and description is not None: + successes.append(description) + continue + + if action == "quit": + parser.exit(0) + elif action == "next": + continue + + if error is not None and mode == "first-error": + break + + assert not old_path.exists() + + return successes, failures + finally: + project.close() + shutil.rmtree(parallel_tree_for_rope, ignore_errors=True) + + +def rename_external( + project_root, + old_name, + new_name, + pattern, + mode, + quiet, + exclude_globs: List[str], +): + if not quiet: + print( + f"rename_external({repr(project_root)}, {repr(old_name)}, {repr(new_name)}, {repr(pattern)})" + ) + project = Project(project_root) + + old_parts = old_name.split(".") + old_path = Path(*old_parts) + + resources = project.get_python_files() + resources = [p for p in resources if fnmatch.fnmatchcase(p.path, pattern)] + resources = [p for p in resources if not fnmatch_any_glob(p.path, exclude_globs)] + + fake_package = tempfile.mkdtemp() + sys.path.append(fake_package) + + os.makedirs(Path(fake_package) / old_path) + (Path(fake_package) / old_path / "__init__.py").touch() + + old_mod = project.find_module(old_name) + + changes = Rename(project, old_mod).get_changes(new_name, resources=resources) + + try: + return apply_changes(project, changes, mode=mode) + finally: + project.close() + shutil.rmtree(fake_package, ignore_errors=True) + sys.path.remove(fake_package) + + +def main(): + args = parser.parse_args() + + global CATCH_ERRORS + CATCH_ERRORS = not args.dont_catch + + exclude_globs = args.exclude_glob or [] + + successes, failures = [], [] + + for old_name, new_name, pattern in args.rename_external or []: + action, description, error, tb = rename_external( + args.repo_root, + old_name, + new_name, + pattern, + mode=args.mode, + quiet=args.quiet, + exclude_globs=exclude_globs, + ) + if error is not None: + failures.append((description, error if args.quiet else tb)) + if error is None and description is not None: + successes.append(description) + elif action == "next": + continue + elif action == "quit": + parser.exit(0) + + convert_to_packages(args.repo_root, exclude_globs=exclude_globs) + + _successes, _failures = prefix_modules( + args.repo_root, + args.prefix, + mode=args.mode, + verbose=args.verbose, + exclude_globs=exclude_globs, + ) + successes.extend(_successes) + failures.extend(_failures) + + if args.mode != "interactive" and not args.quiet: + for description in successes: + print("Successfully applied the patch:") + print(indent(description)) + + for description, e in failures: + print("Failed to apply the patch:") + print(indent(description)) + print(f"The error was: ({type(e).__name__}) {e}") + + if not args.quiet and failures: + print(f"Observed the total of {len(failures)} failures") + + if failures: + parser.exit(1) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 4fba96e..976dddb 100644 --- a/setup.py +++ b/setup.py @@ -16,7 +16,8 @@ entry_points={ "console_scripts": [ "python-package-init = nixpkgs_pytools.python_package_init:main", - "python-rewrite-imports = nixpkgs_pytools.import_rewrite:main" + "python-rewrite-imports = nixpkgs_pytools.import_rewrite:main", + "python-prefix-modules = nixpkgs_pytools.prefix_modules:main", ] }, classifiers=[ diff --git a/tests/test_prefix_modules.py b/tests/test_prefix_modules.py new file mode 100644 index 0000000..9ba8633 --- /dev/null +++ b/tests/test_prefix_modules.py @@ -0,0 +1,118 @@ +import os +import shutil +import sys +import tempfile +from contextlib import contextmanager +from pathlib import Path +from subprocess import check_output + +import pytest +from rope.base.fscommands import subprocess + +from nixpkgs_pytools.prefix_modules import prefix_modules, rename_external + + +@contextmanager +def temp_tree(): + path = tempfile.mkdtemp() + try: + yield Path(path) + finally: + shutil.rmtree(path, ignore_errors=True) + + +@contextmanager +def extended_sys_path(path): + sys.path.append(path) + try: + yield + finally: + # May remove from the wrong index, but close enough + sys.path.remove(path) + + +@contextmanager +def remember_sys_path(): + path = list(sys.path) + try: + yield + finally: + sys.path.clear() + sys.path.extend(path) + + +@pytest.fixture(scope="function") +def some_package(): + with temp_tree() as path: + (path / "utils").mkdir() + (path / "utils" / "__init__.py").touch() + (path / "models.py").write_text("import utils") + + env = { + "PYTHONPATH": ":".join([path.as_posix(), *sys.path]), + } + + with remember_sys_path(): + prefix_modules( + path, + "some_project", + mode="first-error", + verbose=True, + exclude_globs=[], + ) + + yield (path, env) + + +def test_package_moved(some_package): + _, env = some_package + + subprocess.run( + [sys.executable, "-c", "import some_project.utils"], + env=env, + check=True, + ) + + +def test_module_moved(some_package): + path, _ = some_package + + assert (path / "some_project" / "models.py").exists() + + +def test_imports_rewritten(some_package): + _, env = some_package + + subprocess.run( + [sys.executable, "-c", "import some_project.models"], + env=env, + check=True, + ) + + +def test_old_module_doesnt_exist(some_package): + _, env = some_package + with pytest.raises(subprocess.CalledProcessError): + subprocess.run( + [sys.executable, "-c", "import utils"], + env=env, + check=True, + ) + + +def test_rename_external(some_package): + path, _ = some_package + + (path / "thing.py").write_text("import non_existence") + rename_external( + project_root=path, + old_name="non_existence", + new_name="other_project.non_existence", + pattern="**", + mode="first-error", + quiet=False, + exclude_globs=[], + ) + + with open(path / "thing.py", "r") as f: + assert f.read().strip() == "import other_project.non_existence"