如何在不修改Numpy源码的情况下为np.savez_compressed添加压缩级别？

问题背景 #

我需要修改np.savez_compressed内部使用的ZIP压缩级别，Numpy的GitHub上已有相关功能提案但尚未实现。目前有两个可选方案：

• 修改源码文件/numpy/lib/npyio.py，将zipf = zipfile_factory(file, mode="w", compression=compression)替换为zipf = zipfile_factory(file, mode="w", compression=compression, compresslevel=compresslevel)，但每次重装或升级Numpy都要重复修改，并非最优解。
• 对_savez函数做猴子补丁，但尝试时出现ValueError: seek of closed file错误。

错误原因 #

你编写的补丁函数存在以下问题：

1. 缺少os_fspath和format的导入，这两个是Numpy内部依赖的对象，直接使用会触发未定义错误。
2. 手动调用zipf.close()会导致文件提前关闭，后续Numpy内部逻辑依赖该文件对象时就会触发seek操作的错误。
3. 没有正确处理文件对象的上下文管理逻辑。

正确的猴子补丁实现 #

下面是修正后的代码，核心思路是复用Numpy原函数的逻辑，仅修改压缩级别相关的关键部分，避免兼容性问题：

import numpy as np
from numpy.lib import format
from numpy.compat import os_fspath

# 保存原函数，方便后续恢复（可选）
original_savez = np.lib.npyio._savez

def patched_savez(file, args, kwds, compress, allow_pickle=True, pickle_kwargs=None):
    import zipfile
    # 复用原函数的前置处理逻辑
    if not hasattr(file, 'write'):
        file = os_fspath(file)
        if not file.endswith('.npz'):
            file = file + '.npz'
    namedict = kwds
    for i, val in enumerate(args):
        key = 'arr_%d' % i
        if key in namedict:
            raise ValueError(f"Cannot use un-named variables and keyword {key}")
        namedict[key] = val
    if compress:
        compression = zipfile.ZIP_DEFLATED
    else:
        compression = zipfile.ZIP_STORED
    
    # 关键修改：添加自定义compresslevel参数，范围1-9（1最快，9压缩率最高）
    zipf = np.lib.npyio.zipfile_factory(
        file, 
        mode="w", 
        compression=compression, 
        compresslevel=2  # 可根据需求调整数值
    )
    
    try:
        for key, val in namedict.items():
            fname = key + '.npy'
            val = np.asanyarray(val)
            with zipf.open(fname, 'w', force_zip64=True) as fid:
                format.write_array(fid, val, allow_pickle=allow_pickle, pickle_kwargs=pickle_kwargs)
    finally:
        zipf.close()

# 替换原函数
np.lib.npyio._savez = patched_savez

# 测试代码
x = np.array([1, 2, 3, 4])
with open("test.npz", "wb") as f:
    np.savez_compressed(f, x=x)

# 若需要恢复原函数，执行以下代码：
# np.lib.npyio._savez = original_savez

核心说明 #

• 导入了原函数依赖的format和os_fspath，避免未定义错误。
• 使用try...finally确保ZipFile对象正确关闭，同时避免提前关闭导致的文件操作异常。
• 仅修改添加compresslevel的关键行，其余逻辑与原函数保持一致，降低版本兼容性风险。
• 压缩级别范围为1-9，可根据速度/压缩率需求自由调整。

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