77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import ast
|
|
from pathlib import Path
|
|
|
|
|
|
MARKER = " # PRESERVATION_METHODS\n"
|
|
|
|
|
|
def method_span(node: ast.FunctionDef | ast.AsyncFunctionDef) -> tuple[int, int]:
|
|
start = min((decorator.lineno for decorator in node.decorator_list), default=node.lineno)
|
|
if node.end_lineno is None:
|
|
raise ValueError(f"Missing end position for {node.name}")
|
|
return start - 1, node.end_lineno
|
|
|
|
|
|
def move_methods(
|
|
source_path: Path,
|
|
class_name: str,
|
|
target_path: Path,
|
|
method_names: list[str],
|
|
) -> None:
|
|
source = source_path.read_text(encoding="utf-8")
|
|
tree = ast.parse(source, filename=str(source_path))
|
|
owner = next(
|
|
(
|
|
node
|
|
for node in tree.body
|
|
if isinstance(node, ast.ClassDef) and node.name == class_name
|
|
),
|
|
None,
|
|
)
|
|
if owner is None:
|
|
raise ValueError(f"Class not found: {class_name}")
|
|
|
|
methods = {
|
|
node.name: node
|
|
for node in owner.body
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
|
}
|
|
missing = [name for name in method_names if name not in methods]
|
|
if missing:
|
|
raise ValueError(f"Methods not found in {class_name}: {', '.join(missing)}")
|
|
|
|
lines = source.splitlines(keepends=True)
|
|
ordered = sorted((methods[name] for name in method_names), key=lambda node: node.lineno)
|
|
blocks = ["".join(lines[start:end]).rstrip() for start, end in map(method_span, ordered)]
|
|
|
|
for start, end in sorted(map(method_span, ordered), reverse=True):
|
|
del lines[start:end]
|
|
while start < len(lines) - 1 and lines[start] == "\n" and lines[start + 1] == "\n":
|
|
del lines[start]
|
|
|
|
target = target_path.read_text(encoding="utf-8")
|
|
if target.count(MARKER) != 1:
|
|
raise ValueError(f"Target must contain exactly one method marker: {target_path}")
|
|
target = target.replace(MARKER, "\n\n".join(blocks) + "\n")
|
|
|
|
source_path.write_text("".join(lines), encoding="utf-8")
|
|
target_path.write_text(target, encoding="utf-8")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Mechanically move class methods between modules")
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--class-name", required=True)
|
|
parser.add_argument("--target", type=Path, required=True)
|
|
parser.add_argument("methods", nargs="+")
|
|
args = parser.parse_args()
|
|
move_methods(args.source, args.class_name, args.target, args.methods)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|