如何索引CSV文件行并以日期为参数查询对应收盘价？

原始数据与代码 #

CSV样本

Date       Open    High    Low    Close
9/2/2021   34.05   40.34   33.03  36.7
9/3/2021   35.9    41.98   34.9   36.89

现有代码

import csv

def StockMarket():
    while True:
       command = input('$')
       if command.lower() == 'quit':
           break
       elif command.lower() == 'readfiles':
           mrna_data,pfe_data = ReadFiles('MRNA.csv.numbers','PFE.1.csv')
    #elif command.lower() = 'pricesondate'
        

def ReadFiles(MRNA,PFE):
    file = open(MRNA,'r', newline = '')
    reader = csv.reader(file,delimiter ='\t')
    next(reader)
    mrna_data = []
    for row in reader:
        #index rows only need row 0 and 4th row
        mrna_data.append(row)
    reader.close()

    file = open(PFE,'r', newline = '')
    reader = csv.reader(file,delimiter ='\t')
    next(reader)
    pfe_data = []
    for row in reader:
        pfe_data.append(row)
    reader.close()

    return mrna_data,pfe_data

修改后的完整代码 #

import csv
from datetime import datetime

def StockMarket():
    mrna_data, pfe_data = None, None
    while True:
        command = input('$')
        if command.lower() == 'quit':
            break
        elif command.lower() == 'readfiles':
            mrna_data, pfe_data = ReadFiles('MRNA.csv.numbers','PFE.1.csv')
        elif command.lower() == 'pricesondate':
            target_date = input("请输入YYYY-MM-DD格式的日期: ")
            mrna_close = get_close_by_date(mrna_data, target_date)
            pfe_close = get_close_by_date(pfe_data, target_date)
            print(f"MRNA收盘价: {mrna_close}, PFE收盘价: {pfe_close}")

def ReadFiles(mrna_path, pfe_path):
    # 处理MRNA文件，提取日期（格式化后）和收盘价
    mrna_dict = {}
    with open(mrna_path, 'r', newline='') as file:
        reader = csv.reader(file, delimiter='\t')
        next(reader)  # 跳过表头行
        for row in reader:
            if len(row) >= 5:  # 避免列数不足的报错
                raw_date = row[0]
                # 将原始日期（如9/2/2021）转为YYYY-MM-DD格式
                formatted_date = datetime.strptime(raw_date, "%m/%d/%Y").strftime("%Y-%m-%d")
                close_price = float(row[4])
                mrna_dict[formatted_date] = close_price
    
    # 处理PFE文件，逻辑同上
    pfe_dict = {}
    with open(pfe_path, 'r', newline='') as file:
        reader = csv.reader(file, delimiter='\t')
        next(reader)
        for row in reader:
            if len(row) >= 5:
                raw_date = row[0]
                formatted_date = datetime.strptime(raw_date, "%m/%d/%Y").strftime("%Y-%m-%d")
                close_price = float(row[4])
                pfe_dict[formatted_date] = close_price
    
    return mrna_dict, pfe_dict

def get_close_by_date(stock_data, target_date):
    # 按日期查询收盘价，不存在则返回提示信息
    return stock_data.get(target_date, "无对应日期数据")

关键改进点 #

• 数据结构优化：用字典存储（键为YYYY-MM-DD格式日期，值为收盘价），查询效率大幅提升，避免遍历列表
• 日期格式统一：借助datetime模块把原始日期转为标准格式，确保和输入参数匹配
• 文件操作规范：使用with语句自动管理文件资源，无需手动调用close()
• 独立查询函数：get_close_by_date专门处理日期查询逻辑，符合单一职责原则
• 容错处理：检查行的列数，避免索引越界报错

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