git-split.py

#!/usr/bin/env python3
import re
import sys
import subprocess
import argparse
import logging
import random
import string
from collections import defaultdict
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Union

from intervaltree import IntervalTree, Interval
from tabulate import tabulate

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


def run_git_command(command: str, dry_run: bool = False) -> str:
    """
    Run a git command using subprocess and return the output. Exit if an error occurs.

    :param command: Git command to run
    :param dry_run: If True, only print the command without executing it
    :return: Output of the command
    """
    logger.debug(f"Running command: {command}")
    if dry_run:
        print(f"[DRY RUN] {command}")
        return ""
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"Error running command: {command}")
        print(result.stderr)
        sys.exit(1)
    return result.stdout.strip()


def parse_range(range_str: str) -> List[Tuple[Optional[int], Optional[int]]]:
    """
    Parse a range string into a list of start and end tuples.

    :param range_str: Range string (e.g., '10-20,30-,50-60')
    :return: List of tuples with start and end ranges

    >>> parse_range('10-20,30-,50-60')
    [(10, 20), (30, None), (50, 60)]

    >>> parse_range('1-5')
    [(1, 5)]

    >>> parse_range('10-')
    [(10, None)]

    >>> parse_range('-15')
    [(None, 15)]

    >>> parse_range('5-10,15-20')
    [(5, 10), (15, 20)]

    >>> parse_range('-')
    [(None, None)]

    >>> parse_range('1-1,2-2,3-3')
    [(1, 1), (2, 2), (3, 3)]

    >>> parse_range('5-10,20-,-')
    [(5, 10), (20, None), (None, None)]

    >>> parse_range('')
    Traceback (most recent call last):
        ...
    SystemExit: 1

    >>> parse_range('abc')
    Traceback (most recent call last):
        ...
    SystemExit: 1
    """
    ranges = range_str.split(',')
    parsed_ranges = []
    for r in ranges:
        if '-' in r:
            start, end = r.split('-')
            start = int(start) if start else None
            end = int(end) if end else None
            parsed_ranges.append((start, end))
        else:
            print(f"Invalid range format '{r}'. Use N-M, N-, or -M.")
            sys.exit(1)
    return parsed_ranges


def get_current_branch(dry_run: bool = False) -> str:
    """
    Get the name of the current Git branch.

    :return: Current branch name
    """
    return run_git_command("git rev-parse --abbrev-ref HEAD", dry_run)


def generate_random_branch_name() -> str:
    """
    Generate a random branch name.

    :return: Random branch name
    """
    return 'temp-split-' + random_suffix()


def random_suffix() -> str:
    """
    Generate a random suffix of 8 characters.

    :return: Random suffix
    """
    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))


def move_file_and_commit(branch_name: str, old_file: str, new_file: str, commit_message: str,
                         dry_run: bool = False) -> None:
    """
    Move a file in Git, commit the change, and return to the previous branch.

    :param branch_name: Name of the branch to move the file in
    :param old_file: Path to the old file
    :param new_file: Path to the new file
    :param commit_message: Commit message
    :param dry_run: If True, only print the commands without executing them
    """
    run_git_command(f"git checkout -b {branch_name}", dry_run)
    run_git_command(f"git mv {old_file} {new_file}", dry_run)
    run_git_command(f"git add {new_file}", dry_run)
    run_git_command(f"git commit -m '{commit_message}'", dry_run)
    run_git_command("git checkout -", dry_run)


def show_change_digest(
        changes: List[Tuple[str, str, List[str], List[str], List[Tuple[Optional[int], Optional[int]]]]]) -> None:
    """
    Show a digest of the changes to be made.

    :param changes: List of changes
    """
    print("\nChange Digest:")
    table_data = []
    for interval in sorted(changes):
        d = interval.data
        table_data.append([d['origin'], interval.begin, interval.end, d['target']])

    headers = ["Origin", "From", "To", "Spin-off"]
    print(tabulate(table_data, headers, tablefmt="plain"))


def show_preview(changes: List[Dict[str, Dict[str, Union[str, List[str], IntervalTree]]]]) -> None:
    """
    Show a preview of the changes to be made.

    :param changes: List of dictionaries containing change details
    """
    for change in changes:
        for filename, contents in change.items():
            print(f"======= {filename} ==========")
            extracted_lines = contents['lines']
            if len(extracted_lines) > 7:
                print("".join(extracted_lines[:3]))
                print("...")
                print("".join(extracted_lines[-3:]))
            else:
                print("".join(extracted_lines))


def read_file_lines(file_path: str) -> List[str]:
    """
    Read lines from a file.

    :param file_path: Path to the file to read
    :return: List of lines from the file
    """
    with open(file_path, 'r') as file:
        return file.readlines()


def detect_split_blocks(file_path: Union[str, Path]) -> IntervalTree:
    """
    Detect split blocks in the given lines and return the line ranges in a comma-separated format.
    The script looks for special comment markers to identify blocks that should be extracted:

    SPLIT:EXTRACT:<new_filename>
    ... lines to extract ...
    SPLIT:ENDBLOCK

    Multiple blocks can be marked for extraction to different files. Lines outside marked blocks
    remain in the original file.

    :param file_path: File to scan for split blocks
    :return: IntervalTree with the detected split blocks

    Example file marking:
    ```python
    def function1():
        pass

    # SPLIT:EXTRACT:new_module.py
    def function2():
        pass

    def function3():
        pass
    # SPLIT:ENDBLOCK

    def function4():
        pass
    ```
    """
    lines = read_file_lines(file_path)

    blocks = IntervalTree()
    blocks[1:len(lines) + 1] = {"target": file_path, "origin": file_path}  # Belong to original file

    start_pattern = re.compile(r'SPLIT:EXTRACT:\s*(\S+)')
    end_pattern = re.compile(r'SPLIT:ENDBLOCK')
    current_file = None
    start_line = None

    for idx, line in enumerate(lines, start=1):
        start_match = start_pattern.search(line)
        end_match = end_pattern.search(line)

        if start_match:
            current_file = start_match.group(1).strip()
            start_line = idx + 1  # Next line after the SPLIT:EXTRACT line
            blocks.chop(idx - 1, idx + 1)  # Remove the instruction
        elif end_match and current_file:
            end_line = idx  # Line before the SPLIT:ENDBLOCK line
            blocks.chop(idx - 1, idx + 1)  # Remove the instruction
            new_interval = Interval(start_line, end_line, {"target": current_file, "origin": file_path})
            blocks.chop(start_line, end_line)
            blocks.append(new_interval)
            current_file = None
            start_line = None

    logger.debug(blocks)
    return blocks


def process_file(file_path: str, blocks: IntervalTree) -> Dict[str, List[str]]:
    """
    Process the given file and return a dictionary of <file name> -> <contents>.

    :param file_path: Path to the file to process
    :return: Dictionary of file name to list of lines
    """

    content_dict = defaultdict(dict)

    for i, line in enumerate(read_file_lines(file_path)):
        # Find the block (interval) that contains the current line index
        overlapping_intervals = blocks[i + 1]
        if not overlapping_intervals:
            logger.debug(f"{i + 1}: -")
            continue

        data = list(overlapping_intervals)[0].data
        target_file = data['target']  # Assuming one interval per line

        logger.debug(f"{i + 1}: {target_file}: {line.rstrip()}")
        if target_file:
            cd = content_dict[target_file]
            if 'lines' not in cd:
                cd['lines'] = list()
                cd['source'] = file_path
                cd['block'] = overlapping_intervals

            content_dict[target_file]['lines'].append(line)

    return content_dict


def blocks_from_ranges(ranges: List[Tuple[Optional[int], Optional[int]]], target_file: str) -> IntervalTree:
    """
    Create an IntervalTree from a list of line ranges.

    :param ranges: List of tuples with start and end line numbers
    :param target_file: The target file for the intervals
    :return: IntervalTree with the specified ranges
    """
    blocks = IntervalTree()
    for start, end in ranges:
        if start is None and end is None:
            continue
        start = start if start is not None else 1
        end = end if end is not None else sys.maxsize
        blocks[start:end] = {"target": target_file}
    return blocks


def write_lines(filename: str, lines: List[str], dry_run: bool = False) -> None:
    """
    Write lines to a specified file.

    :param filename: Path to the file
    :param lines: List of lines to write
    :param dry_run: If True, only print the action without writing
    """
    if dry_run:
        print(f"Would write {len(lines)} lines to {filename}")
        return
    with open(filename, 'w') as file:
        file.writelines(lines)
    run_git_command(f"git add {filename}", dry_run)


def main() -> None:
    """
    Split a file into multiple files while preserving Git history.

    This script provides two methods for splitting files:

    1. Using inline markers in the source file:
       Add special comments to mark blocks that should be extracted to new files:

       ```python
       # Original content stays in the file
       def function1():
           pass

       # SPLIT:EXTRACT:new_module.py
       def function2():
           pass

       def function3():
           pass
       # SPLIT:ENDBLOCK

       # More original content
       def function4():
           pass
       ```

       Then run: python3 split.py source_file.py

    2. Using command-line line ranges:
       Specify the source file, line ranges, and target file:
       python3 split.py "source.py:10-20,30-40:new_file.py"

       Line range formats:
       - N-M: Extract lines N through M
       - N-: Extract lines N through end of file
       - -M: Extract lines from start through M
       - -: Extract entire file
       - Multiple ranges: "10-20,30-40,50-"

    The script will:
    1. Parse the file and identify blocks to extract
    2. Show a preview of changes (with -p flag)
    3. Create temporary Git branches to preserve history
    4. Move and split files while maintaining Git blame/history
    5. Clean up temporary branches

    Examples:
        # Split using inline markers:
        python3 split.py source_file.py

        # Split using line ranges:
        python3 split.py "source.py:10-20:new_file.py"

        # Preview changes:
        python3 split.py -p source_file.py

        # Multiple files:
        python3 split.py file1.py "file2.py:1-10:new.py"

        # Custom commit message:
        python3 split.py -m "Refactor: Split utility functions" source.py

    The script requires a clean git working directory and will create temporary
    branches during the process. Use --dry-run to preview git operations.
    """
    parser = argparse.ArgumentParser(
        description='Split a file into multiple files while preserving line history.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=main.__doc__)
    parser.add_argument('file_changes', nargs='+',
                        help='List of files to process or specific changes in format old_file:range:new_file')
    parser.add_argument('-m', '--message',
                        help="Final commit message")
    parser.add_argument('-y', '--accept-all', action='store_true',
                        help="Accept all changes without confirmation")
    parser.add_argument('-p', '--preview', action='store_true',
                        help="Show preview of changes")
    parser.add_argument('-n', '--dry-run', action='store_true',
                        help="Perform a dry run without making any changes")

    args = parser.parse_args()

    file_changes = args.file_changes
    final_message = args.message if args.message else "Splitting files"
    dry_run = args.dry_run

    changes = []

    for change in file_changes:
        if re.match(r'^.+?:\d*-\d*:.+?$', change):
            logger.debug(f"Manual mode detected for change: {change}")
            old_file, range_str, new_file = change.split(':')
            line_ranges = parse_range(range_str)
            blocks = blocks_from_ranges(line_ranges, new_file)
        else:
            logger.debug(f"Autodetect mode detected for file: {change}")
            old_file = change
            blocks = detect_split_blocks(old_file)

        content_dict = process_file(old_file, blocks)
        changes.append(content_dict)

    # Show change digest
    show_change_digest(blocks)

    # Show preview if -p is provided
    if args.preview:
        show_preview(changes)

    # Wait for confirmation if -y is not provided
    if not args.accept_all:
        confirm = input("\nDo you want to proceed with these changes? (yes/no): ").strip().lower()
        if confirm != 'yes':
            print("Aborting changes.")
            sys.exit(0)

    # Ensure working directory is clean, not influenced by dry-run
    run_git_command(
        "git diff --quiet || (echo 'Refusing to work on a dirty tree! Please, clean up to continue.' && exit 1)")

    # Get current branch, not influenced by dry-run
    original_branch = get_current_branch()

    # Generate random final branch name
    final_branch = generate_random_branch_name()

    run_git_command(f"git checkout -b {final_branch}", dry_run)

    branches_to_merge = list()
    originals = dict()

    # Process each change
    for change in changes:

        for filename, contents in change.items():
            print(f"Processing file: {filename}...")
            logger.debug(contents.keys())
            logger.debug(contents['lines'])

            original = contents['source']
            if filename == original:
                print(f"Special treatment for file {filename}")
                if original in originals:
                    # We've done with this original file; only once
                    print(f"Original file {filename} was already processed")
                    continue
                filename = f"{filename}.{random_suffix()}"
                originals[original] = filename

            # Generate random branch names
            branch = generate_random_branch_name()
            branches_to_merge.append(branch)

            move_file_and_commit(branch, original, filename,
                                 f"{final_message}: Move {original} to {filename}", dry_run)

        all_branches = ' '.join(branches_to_merge)

        # merge all branches 
        run_git_command(f"git merge {all_branches} -m '{final_message}: Intermediate merge'", dry_run)

        # Move temp file back to old file
        for original, filename in originals.items():
            run_git_command(f"git mv {filename} {original}", dry_run)
            run_git_command(f"git add {original}", dry_run)

        run_git_command(f"git commit -m '{final_message}: Move temp files back to original names'", dry_run)

        # Write extracted lines and remaining lines to respective files
        for filename, contents in change.items():
            lines = contents['lines']
            write_lines(filename, lines, dry_run)

        run_git_command(f"git commit -m '{final_message}: Post-merge edit'", dry_run)

    run_git_command(f"git checkout {original_branch}", dry_run)

    # Merge final branch into original branch
    run_git_command(f"git merge --no-ff {final_branch} -m '{final_message}'", dry_run)

    # Delete the temporary branches
    run_git_command(f"git branch -d {all_branches}", dry_run)
    run_git_command(f"git branch -d {final_branch}", dry_run)

    print("Split completed successfully." if not dry_run else "Dry run completed successfully.")


if __name__ == "__main__":
    main()