Function sanitize

src/sanitize.py:343–431

(
    old_code: str,
    entry_point: str,
    rm_prefix_lines: Optional[str] = None,
    eofs: List = None,
    codegen: bool = False,
    global_code: bool = False,
    chat: bool = False
)

def sanitize(
    old_code: str,
    entry_point: str,
    rm_prefix_lines: Optional[str] = None,
    eofs: List = None,
    codegen: bool = False,
    global_code: bool = False,
    chat: bool = False
):
    new_code = old_code.replace("\r\n", "\n").replace("\\_", "_").replace("if __name__", "if 1 or __name__")
    if codegen:
        new_code = remove_space_for_codegen(new_code)
    if global_code and "```" in new_code:
        if not chat:
            new_code = new_code.split("```")[0]
            return new_code.strip()
        else:
            if len(new_code.split("```python\n")) > 1:
                new_code = new_code.split("```python\n")[1]
            elif len(new_code.split("```")) > 1:
                new_code = new_code.split("```")[1]
            new_code = new_code.split("```")[0]
            
            return new_code.strip()



    if new_code.endswith("```"):
        new_code = new_code[:-3]
    if rm_prefix_lines is not None:
        new_code = "\n".join(
            [
                line
                for line in old_code.splitlines()
                if not line.startswith(rm_prefix_lines)
            ]
        )

    new_code = "\n" + new_code
    def_left = "def " + entry_point

    # basic handling of chat output
    new_code = new_code.replace("\n```python\n", "\n```\n")
    if def_left in new_code:
        for chunk in new_code.split("\n```\n"):
            if def_left in chunk:
                new_code = chunk
                break
    else:
        new_code = new_code.split("```")[0]

    if codegen:
        for chunk in new_code.split("\"\"\""):
            if def_left in chunk:
                new_code = chunk
                break
    
    chunks = [chunk for chunk in re.split(f"{def_left}\s*\(", new_code)]