如何使用argparse实现终端指定DataFrame列的筛选功能？

问题背景 #

现有代码读取CSV文件转为DataFrame，希望通过终端命令./main.py --fields column1,column2指定要展示的列，尝试使用argparse但未成功，同时现有完整代码存在多处逻辑问题导致功能失效。

原尝试的argparse代码存在参数名不匹配问题：

parser = argparse.ArgumentParser(description='Process rows to display.')
parser.add_argument('--fields', type=str, help='Rows to display')

args = parser.parse_args()
select = args.select.split(',') if args.select else None

这里参数定义为--fields，但引用时错误使用args.select，导致无法获取输入字段。

DataFrame示例数据：

date                 column1  column2     column3   column4 
0   2022-01-01          1     2.550000  Unknown  facebook
1   2022-01-01          5     2.470000  Unknown  facebook
2   2022-01-01          2     2.620000  Unknown  facebook
3   2022-01-01          3     2.480000  Unknown  facebook
4   2022-01-01          2     2.440000  Unknown  facebook

修正后的完整代码 #

import pandas as pd
import requests
from io import StringIO
import hashlib
import time
import argparse
import sys


def get_file(selected_fields=None):
    url = 'https://drive.google.com/file/d/1zLdEcpzCp357s3Rse112Lch9EMUWzMLE/view?usp=sharing'
    file_id = url.split('/')[-2]
    dwn_url = 'https://drive.google.com/uc?export=download&id=' + file_id
    url2 = requests.get(dwn_url).text
    csv_raw = StringIO(url2)
    df = pd.read_csv(csv_raw)
    df.drop(5, inplace=True)
    data_json = df.to_json(orient='records')

    # 处理指定字段筛选
    if selected_fields:
        selected_fields = [field.strip() for field in selected_fields]
        valid_fields = [field for field in selected_fields if field in df.columns]
        if not valid_fields:
            print("指定字段均不存在，将展示所有列", file=sys.stderr)
        else:
            df = df[valid_fields]
    
    return df, data_json


def setup_display_options():
    pd.options.display.width = None
    pd.options.display.max_columns = None
    pd.options.display.max_colwidth = None


def calculate_hash(df):
    df_str = df.to_string().encode()
    return hashlib.sha256(df_str).hexdigest()


def monitor_file(interval, selected_fields=None):
    previous_hash = None
    setup_display_options()
    while True:
        df, data_json = get_file(selected_fields)
        current_hash = calculate_hash(df)
        if previous_hash != current_hash:
            print("文件已更新")
            print(df.head(50))
            print(data_json)
        previous_hash = current_hash
        time.sleep(interval)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='指定要展示的DataFrame列')
    parser.add_argument('--fields', type=str, 
                        help='要展示的列，用逗号分隔，例如：--fields date,column1,column2')
    parser.add_argument('--interval', type=int, default=120,
                        help='文件监控间隔（秒），默认120秒')
    
    args = parser.parse_args()
    
    selected_fields = args.fields.split(',') if args.fields else None
    monitor_file(args.interval, selected_fields)

关键修改说明 #

• 修正argparse参数错误：将args.select改为args.fields，匹配定义的参数名，确保能正确获取用户输入的字段。
• 字段输入优化：对输入字段做去空格处理，避免因column1, column2这种带空格的格式导致列名匹配失败；同时增加字段有效性校验，不存在的字段会给出提示。
• 修复函数返回值：原get_file未返回df和data_json，导致monitor_file赋值失败，修改后添加返回逻辑。
• 优化显示设置：将重复的显示参数设置整理为统一函数，减少冗余代码。
• 扩展命令行参数：新增--interval参数，允许用户自定义监控间隔。
• 移除错误的sys.argv处理：删除原代码中直接读取sys.argv的逻辑，统一使用argparse解析参数，保证逻辑一致性。

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