基于argparse参数动态调整设备信息表格展示的问题求助

我来帮你搞定这个动态表格展示的问题！先梳理下代码里的核心问题，再给你修正后的完整实现：

问题根源 #

你的代码目前存在几个关键问题：

• 无条件打印主表格和LLDP表格，不管用户是否传入对应参数
• --outfile的处理逻辑完全失效，没有实现写入文件的功能
• --inventory参数的默认值设为列表会导致文件存在性判断出错
• 主表格即使没有用户指定的额外列（仅Hostname）也会被打印

修正后的完整代码 #

import argparse
import os
import sys
import yaml
from prettytable import PrettyTable

def generate_data(inventory):
    # 模拟从设备获取数据的逻辑，替换为你的实际设备数据获取代码即可
    mock_data = {
        "csr1": {
            "os_version": "16.6",
            "serial_number": "9KIBQAQ3OPE",
            "lldp": [
                {"device-id": "csr2.com", "local-interface": "Gi1", "connecting-interface": "Gi1"}
            ]
        },
        "csr2": {
            "os_version": "16.6",
            "serial_number": "9KIBQAQ3OPE",
            "lldp": [
                {"device-id": "csr1.com", "local-interface": "Gi1", "connecting-interface": "Gi1"}
            ]
        }
    }
    return mock_data

def generate_table(inventory):
    args = arguments()
    data = generate_data(inventory)
    main_headers = ['os_version', 'serial_number']
    main_header = []
    output_content = []

    # 收集用户指定的主表格列
    for arg in vars(args):
        if arg in main_headers and getattr(args, arg):
            main_header.append(arg)

    # 生成主表格（仅当用户指定了至少一个-s/-v参数时）
    if main_header:
        main_header.insert(0, 'Hostname')
        main_table = PrettyTable()
        main_table.field_names = main_header
        for hostname, details in data.items():
            row = [hostname]
            for column in main_table.field_names[1:]:
                row.append(details[column])
            main_table.add_row(row)
        output_content.append(str(main_table))

    # 生成LLDP表格（仅当用户指定--lldp参数时）
    if args.lldp:
        lldp_table = PrettyTable()
        lldp_table.field_names = ['Hostname', 'Neighbor', 'Local Interface', 'Neighbor Interface']
        for hostname, details in data.items():
            for lldp_data in details['lldp']:
                neighbor = lldp_data['device-id']
                local_int = lldp_data['local-interface']
                neigh_int = lldp_data['connecting-interface']
                lldp_table.add_row([hostname, neighbor, local_int, neigh_int])
        output_content.append(str(lldp_table))

    # 处理输出：打印或写入文件
    final_output = "\n".join(output_content)
    if args.outfile:
        with open('device_info.txt', 'w') as f:
            f.write(final_output)
        print(f"输出已写入 device_info.txt")
    else:
        print(final_output)

def arguments():
    parser = argparse.ArgumentParser(description='Argparse for Training Course.')
    parser.add_argument('-s', '--serial_number', action='store_true', help='Device Serial Numbers')
    parser.add_argument('-v', '--os_version', action='store_true', help='Output Devices OS')
    parser.add_argument('--lldp', action='store_true', help='Output LLDP Data')
    parser.add_argument('--outfile', action='store_true', help='Output to file')
    # 修正默认值为字符串，避免文件存在性判断出错
    parser.add_argument('--inventory', help='Inventory File', default="inventory.yml", required=True)
    args = parser.parse_args()
    return args

def get_inventory(inventory):
    with open(inventory) as fh:
        # 使用safe_load更安全，避免yaml注入风险
        yml_file = yaml.safe_load(fh)
    return yml_file

def main():
    args = arguments()
    if not os.path.isfile(args.inventory):
        sys.exit('Please specify valid, readable YAML file with data')
    inventory = get_inventory(args.inventory)
    generate_table(inventory)

if __name__ == '__main__':
    main()

关键改动说明 #

1. 条件生成表格
   • 主表格：只有当用户指定了-s或-v（即main_header不为空）时才生成
   • LLDP表格：只有当用户传入--lldp参数时才生成
2. 修复--outfile功能
   • 将所有需要输出的内容收集到output_content列表，最后统一处理：要么打印到终端，要么写入文件
3. 修正--inventory参数
   • 默认值从列表["inventory.yml"]改为字符串"inventory.yml"，避免文件存在性判断出错
4. 安全的YAML加载
   • 把yaml.load改为yaml.safe_load，避免潜在的安全风险
5. 移除无效逻辑
   • 删除了原来无效的lldp_header和out_file列表处理逻辑

测试验证 #

现在你可以测试各个场景，都会完全符合你的预期：

• python3 rest5.py --inventory inventory.yml -v：仅显示带os_version的主表格
• python3 rest5.py --inventory inventory.yml --lldp：仅显示LLDP表格
• python3 rest5.py --inventory inventory.yml -s -v --lldp：同时显示主表格和LLDP表格
• python3 rest5.py --inventory inventory.yml -s --outfile：将主表格写入device_info.txt

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