llama_patch.py

import click
import re
import sys
import os
import difflib
import logging

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

class LlamaPatchException(Exception):
    pass

@click.command()
@click.option('--llmpatch', type=click.File('r'), default='-', help='Patch file to apply, defaults to STDIN if not provided.')
@click.option('--out', type=click.File('w'), default='-', help='Output file, defaults to STDOUT if not provided.')
@click.option('--cwd', type=click.Path(), default='.', help='Current working directory, defaults to ".".')
@click.option('--verbose', '-v', is_flag=True, help='Enable verbose logging.')
def apply_patch(llmpatch, out, cwd, verbose):
    if verbose:
        logger.setLevel(logging.DEBUG)
    
    os.chdir(cwd)
    
    patch_content = llmpatch.read()
    patches = parse_patches(patch_content)

    try:
        diff_output = []
        for patch in patches:
            file_path, element_type, element_name, changes = patch
            logger.debug(f"Processing patch for {element_type} {element_name} in {file_path}")
            diff_output.append(generate_diff_for_patch(file_path, element_type, element_name, changes))
        
        if out.name != '<stdout>':
            with open(out.name, 'w') as file:
                file.write('\n'.join(diff_output))
        else:
            sys.stdout.write('\n'.join(diff_output))
    
    except LlamaPatchException as e:
        logger.error(f"Error applying patch: {e}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        sys.exit(1)

def parse_patches(patch_content):
    patch_pattern = re.compile(r'^--- (.+?)\n\?\? (\w+)\s*(\w+)?\n((?:[+\-].*?\n)+)', re.DOTALL | re.MULTILINE)
    matches = patch_pattern.findall(patch_content)
    
    if not matches:
        raise LlamaPatchException("No valid patches found in the provided patch content.")
    
    patches = []
    for match in matches:
        path, element_type, element_name, changes = match
        changes = changes.strip().split('\n')
        patches.append((path, element_type, element_name, changes))
    
    return patches

def generate_diff_for_patch(file_path, element_type, element_name, changes):
    if not os.path.isfile(file_path):
        raise LlamaPatchException(f"File not found: {file_path}")

    with open(file_path, 'r') as file:
        original_code = file.read()

    if element_type in ["function", "def"]:
        updated_code = apply_function_patch(original_code, element_name, changes)
    elif element_type == "class":
        updated_code = apply_class_patch(original_code, element_name, changes)
    elif element_type == "struct":
        updated_code = apply_struct_patch(original_code, element_name, changes)
    elif element_type == "<<":
        updated_code = prepend_code(original_code, changes)
    elif element_type == ">>":
        updated_code = append_code(original_code, changes)
    elif element_type == "call":
        updated_code = apply_call_patch(original_code, element_name, changes)
    else:
        raise LlamaPatchException(f"Unsupported element type: {element_type}")

    original_lines = original_code.splitlines(keepends=True)
    updated_lines = updated_code.splitlines(keepends=True)

    diff = difflib.unified_diff(original_lines, updated_lines, fromfile=file_path, tofile=file_path)
    return ''.join(diff)

def apply_function_patch(original_code, function_name, changes):
    function_pattern = re.compile(rf'def {function_name}\(.*?\):\n((?:\s+.*?\n)*)', re.DOTALL)
    match = function_pattern.search(original_code)
    if not match:
        raise LlamaPatchException(f"Function {function_name} not found")

    start, end = match.span()
    context_code = match.group(1)

    new_code = generate_new_code(context_code, changes)
    updated_code = original_code[:start] + f'def {function_name}(\n{new_code}' + original_code[end:]
    return updated_code

def apply_class_patch(original_code, class_name, changes):
    class_pattern = re.compile(rf'class {class_name}:\n((?:\s+.*?\n)*)', re.DOTALL)
    match = class_pattern.search(original_code)
    if not match:
        raise LlamaPatchException(f"Class {class_name} not found")

    start, end = match.span()
    context_code = match.group(1)

    new_code = generate_new_code(context_code, changes)
    updated_code = original_code[:start] + f'class {class_name}:\n{new_code}' + original_code[end:]
    return updated_code

def apply_struct_patch(original_code, struct_name, changes):
    struct_pattern = re.compile(rf'struct {struct_name} {{\n((?:\s+.*?\n)*)}}', re.DOTALL)
    match = struct_pattern.search(original_code)
    if not match:
        raise LlamaPatchException(f"Struct {struct_name} not found")

    start, end = match.span()
    context_code = match.group(1)

    new_code = generate_new_code(context_code, changes)
    updated_code = original_code[:start] + f'struct {struct_name} {{\n{new_code}}}' + original_code[end:]
    return updated_code

def apply_call_patch(original_code, call_name, changes):
    call_pattern = re.compile(rf'{call_name}\(.*?\)', re.DOTALL)
    matches = list(call_pattern.finditer(original_code))
    if not matches:
        raise LlamaPatchException(f"Function call {call_name} not found")

    updated_code = original_code
    for match in matches:
        start, end = match.span()
        context_code = match.group()

        try:
            new_code = generate_new_code(context_code, changes, call=True)
            updated_code = updated_code[:start] + new_code + updated_code[end:]
        except LlamaPatchException as e:
            raise LlamaPatchException(f"Error applying patch for call {call_name}: {str(e)}")

    return updated_code

def prepend_code(original_code, changes):
    new_code = "\n".join(line[1:] for line in changes if line.startswith('+'))
    updated_code = new_code + "\n" + original_code
    return updated_code

def append_code(original_code, changes):
    new_code = "\n".join(line[1:] for line in changes if line.startswith('+'))
    updated_code = original_code + "\n" + new_code
    return updated_code

def generate_new_code(context_code, changes, call=False):
    context_lines = context_code.split('\n')
    new_code = []

    try:
        for line in changes:
            if line.startswith('-'):
                to_remove = line[1:].strip()
                context_lines = [l for l in context_lines if l.strip() != to_remove]
            elif line.startswith('+'):
                new_code.append(line[1:])
            else:
                new_code.append(line)
    except ValueError as e:
        raise LlamaPatchException(f"Error in changes: {str(e)}\nContext code:\n{context_code}\nChanges:\n{changes}")

    if call:
        return "".join(new_code)
    else:
        return "\n".join(new_code)

if __name__ == "__main__":
    apply_patch()