优化按评分拆分JSON Lines数据的Python代码方案咨询

需求背景 #

现有JSON Lines格式的亚马逊视频评论数据，需要按overall字段的1-5分，将对应reviewText提取到ranking1.txt至ranking5.txt文件中。原代码可正常运行，但存在冗余，需要优化效率与可读性。

原实现代码 #

f1 = open("ranking1.txt", "a")
f2 = open("ranking2.txt", "a")
f3 = open("ranking3.txt", "a")
f4 = open("ranking4.txt", "a")
f5 = open("ranking5.txt", "a")

with open("/content/Amazon_Instant_Video_5.json") as json_file:
  for line in json_file: #runs the loop to extract info
    data = json.loads(line)
    if data['overall'] == 1:
      f1.write(data['reviewText'] + '\n')
    if data['overall'] ==2:
      f2.write(data['reviewText'] + '\n')
    if data['overall'] == 3:
      f3.write(data['reviewText'] + '\n')
    if data['overall'] ==4:
      f4.write(data['reviewText'] + '\n')
    if data['overall'] ==5:
      f5.write(data['reviewText'] + '\n')

优化后的代码 #

基础优化版本 #

# 用字典映射评分到对应的文件句柄
file_handlers = {}
for score in range(1, 6):
    file_handlers[score] = open(f"ranking{score}.txt", "a")

try:
    with open("/content/Amazon_Instant_Video_5.json") as json_file:
        for line in json_file:
            data = json.loads(line)
            # 将浮点型评分转为整数（如2.0→2），避免类型匹配问题
            score = int(data['overall'])
            # 只处理1-5分的有效评论
            if 1 <= score <=5:
                file_handlers[score].write(f"{data['reviewText']}\n")
finally:
    # 确保所有文件句柄被正确关闭
    for fh in file_handlers.values():
        fh.close()

更简洁的上下文管理器版本（Python 3.10+） #

from contextlib import ExitStack

with ExitStack() as stack:
    # 批量打开文件并自动管理关闭
    file_handlers = {
        score: stack.enter_context(open(f"ranking{score}.txt", "a"))
        for score in range(1, 6)
    }
    
    with open("/content/Amazon_Instant_Video_5.json") as json_file:
        for line in json_file:
            data = json.loads(line)
            score = int(data['overall'])
            if 1 <= score <=5:
                file_handlers[score].write(f"{data['reviewText']}\n")

优化点说明 #

• 减少冗余代码：用字典统一管理文件句柄，避免重复定义f1-f5变量，后续扩展评分范围时只需修改循环范围即可。
• 提升匹配效率：原代码中多个独立if会逐个判断，优化后直接通过字典索引定位文件，无需多次条件判断；若保留条件判断，建议用elif替代多个if，减少不必要的判断步骤。
• 增强代码安全性：用try/finally或ExitStack确保所有文件句柄被正确关闭，原代码未显式关闭文件，可能导致资源泄漏。
• 兼容性处理：将浮点型的overall转为整数，避免因类型不一致引发的潜在问题，同时增加范围判断，过滤异常评分值。

内容的提问来源于stack exchange，提问作者Anano Mikadze