一个将代码块统一添加代码行号的脚本

在 LLM Agent 训练中，有时存在需要通过代码行号进行补全的方法。

这个脚本给任意给定的代码行统一添加代码行号。

import json
import re
import argparse

def add_line_numbers_to_input(input_jsonl_path, output_jsonl_path):
    with open(input_jsonl_path, 'r', encoding='utf-8') as fin, \
         open(output_jsonl_path, 'w', encoding='utf-8') as fout:
        
        for line in fin:
            try:
                data = json.loads(line)
                if "input" in data:
                    data["input"] = process_code_block(data["input"])
                fout.write(json.dumps(data, ensure_ascii=False) + '\n')
            except json.JSONDecodeError:
                print(f"JSON 解码错误: {line}")

def process_code_block(code_str):
    # 匹配代码块（包含```c和```， 如果需要匹配其他语言可以之后更改）
    code_block_match = re.search(r'(```c\n)(.*?)(\n```)', code_str, re.DOTALL)

    if not code_block_match:
        print("未找到符合条件的代码块")
        return code_str
    
    # 获取代码块各部分
    start_marker = code_block_match.group(1)  # ```c\n
    code_content = code_block_match.group(2)   # 实际代码内容
    end_marker = code_block_match.group(3)    # \n```

    # 处理行号（保留原始缩进）
    code_lines = code_content.split('\n')
    number_lines = []
    line_num = 1

    for code_line in code_lines:
        number_lines.append(f"{line_num}: {code_line}")
        line_num += 1
    
    # 重新建立代码块
    processed_code = (
        code_str[:code_block_match.start()] +  # 保留代码块前的所有内容
        start_marker +                        # ```c\n
        '\n'.join(number_lines) +           # 带行号的代码
        end_marker +                          # \n```
        code_str[code_block_match.end():]     # 保留代码块后的所有内容
    )

    return processed_code

def main():
    parser = argparse.ArgumentParser(description="Add line numbers to code blocks in a JSONL file.")
    parser.add_argument('-i', '--input', required=True, help='Input JSONL file path')
    parser.add_argument('-o', '--output', required=False, default='output.jsonl', help='Output JSONL filepath, default=output.jsonl')
    args = parser.parse_args()

    add_line_numbers_to_input(args.input, args.output)
    print(f"Processed file saved to {args.output}")


if __name__  == '__main__':
    main()