def extract_program()

in evaluation/trajectory.py [0:0]

• 34 lines of code
• 16 McCabe index (conditional complexity)

def extract_program(text:str=None, trajectory:list=None, last_only=False) -> str:
    assert text is not None or trajectory is not None, "Either text or trajectory should be provided."
    if trajectory is None:
        try:
            trajectory = text_to_trajectory(text)
        except:
            return "raise ValueError('Invalid trajectory')"

    program_list = []
    import_lines = []
    for i, item in enumerate(trajectory):
        if item["role"] == "program":
            cur_program = item["content"]
            if i < len(trajectory) - 1:
                assert trajectory[i+1]["role"] == "output"
                output = trajectory[i+1]["content"].strip()
                if is_execution_success(output):
                    program_list.append(cur_program)
                else:
                    # extract import lines only
                    for line in cur_program.split("\n"):
                        if line.startswith("import") or line.startswith("from"):
                            import_lines.append(line)
            else:
                program_list.append(cur_program)
    # add import lines to the first program
    if len(program_list) == 0:
        program_list.append("")
    if len(import_lines) > 0:
        program_list[0] = "\n".join(import_lines) + "\n" + program_list[0]
    for i, program in enumerate(program_list[:-1]):
        program_list[i] = "\n".join([line for line in program.split("\n") if not line.strip().startswith("print(")])

    if last_only:
        program = program_list[-1]
    else:
        program = "\n".join(program_list)
    return program