Function replace_base_class

src/utils.py:91–128

(code: str, new_class_def: str)

def replace_base_class(code: str, new_class_def: str) -> str:
    lines = code.splitlines(keepends=True)
    class_start = None
    class_end = None

    # Find the start line of class TeachingScene(Scene):
    for i, line in enumerate(lines):
        if re.match(r"^\s*class\s+TeachingScene\s*\(Scene\)\s*:", line):
            class_start = i
            break

    if class_start is not None:
        # Find the end line of the class definition
        # The class ends when a line with the same or less indentation is found
        base_indent = len(lines[class_start]) - len(lines[class_start].lstrip())
        class_end = class_start + 1
        while class_end < len(lines):
            line = lines[class_end]
            # If an empty line or a line with less indentation is found,
            # it means the class definition has ended
            if line.strip() != "" and (len(line) - len(line.lstrip()) <= base_indent):
                break
            class_end += 1

        # Replace the original TeachingScene definition with the new one
        new_block = new_class_def.strip() + "\n\n"
        return "".join(lines[:class_start]) + new_block + "".join(lines[class_end:])
    else:
        # If TeachingScene does not exist, it should be inserted before the first class definition
        for i, line in enumerate(lines):
            if re.match(r"^\s*class\s+\w+", line):
                insert_pos = i
                break
        else:
            insert_pos = 0

        new_block = new_class_def.strip() + "\n\n"
        return "".join(lines[:insert_pos]) + new_block + "".join(lines[insert_pos:])


# Save the program to the.py file