基于argparse的网络流量测量二选一功能实现求助

1. 先搞定argparse参数互斥校验 #

首先要在argparse中设置互斥参数组，强制-t和-n二选一，从根源避免非法输入：

import argparse

parser = argparse.ArgumentParser(description='网络收发数据测量工具')
# 创建互斥参数组，要求用户必须选择其中一个参数
mutex_group = parser.add_mutually_exclusive_group(required=True)
mutex_group.add_argument('-t', '--time', type=int, help='测量时长（单位：秒）')
mutex_group.add_argument('-n', '--num', type=int, help='指定收发的总字节数')
# 可选：自定义缓冲区大小，默认1MB
parser.add_argument('--buffer-size', type=int, default=1024*1024, help='单次收发缓冲区大小，默认1MB')
args = parser.parse_args()

这里用type=int直接完成整数校验，无需额外处理非整数输入的异常。

2. 修正代码逻辑问题 #

你的现有代码存在循环提前终止、未定义变量、仅处理发送等问题，以下是修正后的完整实现：

import time
import socket

# 假设已完成socket初始化（根据实际场景调整连接逻辑）
clientSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
clientSocket.connect(('目标IP', 目标端口))

start_time = time.time()
total_bytes_sent = 0
total_bytes_received = 0
BUFFER_SIZE = args.buffer_size

if args.time is not None:
    duration = args.time
    print(f"测量{duration}秒内的收发数据量...")
    end_time = start_time + duration
    while time.time() < end_time:
        # 发送数据
        send_data = b'\x00' * BUFFER_SIZE
        sent = clientSocket.send(send_data)
        total_bytes_sent += sent
        # 接收服务器回传数据（需配合服务器逻辑，若无需回传可移除）
        recv_data = clientSocket.recv(BUFFER_SIZE)
        total_bytes_received += len(recv_data)
elif args.num is not None:
    target_bytes = args.num
    print(f"测量收发{target_bytes}字节所需时间...")
    while total_bytes_sent + total_bytes_received < target_bytes:
        # 控制单次发送量，避免最后一次超出目标字节数
        remaining = target_bytes - (total_bytes_sent + total_bytes_received)
        send_size = min(BUFFER_SIZE, remaining)
        send_data = b'\x00' * send_size
        sent = clientSocket.send(send_data)
        total_bytes_sent += sent
        # 接收回传数据
        recv_data = clientSocket.recv(send_size)
        total_bytes_received += len(recv_data)

clientSocket.close()

elapsed_time = time.time() - start_time
total_bytes = total_bytes_sent + total_bytes_received
bandwidth = (total_bytes / elapsed_time) * 8 / 1000000  # 转换为Mbps单位

# 输出测量结果
print(f"总收发字节数: {total_bytes} bytes")
print(f"耗时: {elapsed_time:.2f} seconds")
print(f"带宽: {bandwidth:.2f} Mbps")

关键优化点 #

• 双向测量：同时统计发送和接收字节数，符合“收发数据测量”的需求
• 循环逻辑：指定时长时用end_time减少重复计算；指定字节数时用min控制单次发送量，避免超出目标
• 参数合法性：通过argparse互斥组强制二选一，无需额外判断参数是否同时存在

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