使用bz2与csv.writer写入压缩CSV时遇TypeError的解决方法

背景 #

我需要在写入磁盘前压缩CSV文件，因为要在SMP上同时运行约96个进程，不然有限的硬盘空间会被填满（这不是我的系统，别问为啥一台配备104核CPU、256GB内存、8块Tesla显卡的服务器，只给所有用户提供90%以上已占用的2TB共享存储）。单CPU运行要近4年，用96核能缩短到约2周，所以必须用尽可能多的处理器。

类似问题的答案都建议用bz2.open()并设置模式为'wt'，但我没找到把bz2类文件对象和csv.writer()结合的可行方案，这么做似乎没法正常工作。我甚至写了个脚本测试所有可能的写入模式组合，用来复现问题。

注意 #

我不能直接用','.join(row)（这个方法在mode='wt'的bz2对象上能正常工作），因为很多文本字段需要转义换行符、嵌入逗号、嵌入'\x00'字符等。

复现步骤 #

/tmp/test.py代码：

import os
import bz2
import csv
import traceback


tfile = '/tmp/test.csv.bz'
row = ['bc22jtr', 118324, None, 'contran', None, 11.5, 9.23, ]


def perr(err, bmode, fmode=None):
    """Func for printing exception info in a less noisy manner."""
    print(
        f"EXCEPTION: wt.writerow(row) → {type(err).__name__}:"
        f" {err}; bmode='{bmode}', fmode='{fmode}'"
    )
    print((''.join(traceback.format_exception(err)[-2:-1])).strip())
    return True


for fmode in ["w", "wt", "wb"]:
    for bmode in ["w", "wb"]:
        had_err = False
        if os.path.exists(tfile):
            os.remove(tfile)
        fh = open(tfile, fmode)
        try:
            bh = bz2.BZ2File(fh, mode=bmode, compresslevel=9)
        except ValueError as err:
            had_err = perr(err, bmode, fmode)
        wt = csv.writer(fh)  # 核心错误：应绑定到bz2对象bh而非原始文件句柄fh
        try:
            wt.writerow(row)
        except TypeError as err:
            had_err = perr(err, bmode, fmode)
        try:
            bh.close()
        except TypeError as err:
            had_err = perr(err, bmode, fmode)
        if not had_err:
            print(f"WAS OK: bmode={bmode}, fmode={fmode}")  # 原代码prnt为笔误，修正为print
for bmode in ["w", "wb", "wt"]:
    if os.path.exists(tfile):
        os.remove(tfile)
    bh = bz2.open(fh, mode=bmode, compresslevel=9)
    wt = csv.writer(fh)  # 同样错误：绑定到了原始文件句柄fh
    had_err = False
    try:
        wt.writerow(row)
    except TypeError as err:
        had_err = perr(err, bmode)
    try:
        bh.close()
    except TypeError as err:
        had_err = perr(err, bmode)
    if not had_err:
        print(f"WAS OK: bmode={bmode}")
if os.path.exists(tfile):
    os.remove(tfile)

运行输出 #

> python3 /tmp/test.py
EXCEPTION: wt.writerow(row) → TypeError: write() argument must be str, not bytes; bmode='w', fmode='w'
File "/usr/lib/python3.10/bz2.py", line 109, in close
    self._fp.write(self._compressor.flush())
EXCEPTION: wt.writerow(row) → TypeError: write() argument must be str, not bytes; bmode='wb', fmode='w'
File "/usr/lib/python3.10/bz2.py", line 109, in close
    self._fp.write(self._compressor.flush())
EXCEPTION: wt.writerow(row) → TypeError: write() argument must be str, not bytes; bmode='w', fmode='wt'
File "/usr/lib/python3.10/bz2.py", line 109, in close
    self._fp.write(self._compressor.flush())
EXCEPTION: wt.writerow(row) → TypeError: write() argument must be str, not bytes; bmode='wb', fmode='wt'
File "/usr/lib/python3.10/bz2.py", line 109, in close
    self._fp.write(self._compressor.flush())
EXCEPTION: wt.writerow(row) → TypeError: a bytes-like object is required, not 'str'; bmode='w', fmode='wb'
File "/tmp/test.py", line 33, in <module>
    wt.writerow(row)
EXCEPTION: wt.writerow(row) → TypeError: a bytes-like object is required, not 'str'; bmode='wb', fmode='wb'
File "/tmp/test.py", line 33, in <module>
    wt.writerow(row)
EXCEPTION: wt.writerow(row) → TypeError: a bytes-like object is required, not 'str'; bmode='w', fmode='None'
File "/tmp/test.py", line 49, in <module>
    wt.writerow(row)
EXCEPTION: wt.writerow(row) → TypeError: a bytes-like object is required, not 'str'; bmode='wb', fmode='None'
File "/tmp/test.py", line 49, in <module>
    wt.writerow(row)
EXCEPTION: wt.writerow(row) → TypeError: a bytes-like object is required, not 'str'; bmode='wt', fmode='None'
File "/tmp/test.py", line 49, in <module>
    wt.writerow(row)

额外注意 #

第一个循环里没测试bmode='wt'，因为bz2.BZ2File(fh, mode='wt')总会抛出ValueError: Invalid mode: 'wt'异常。

核心问题 #

如何在写入时生成带有正确转义和编码的压缩CSV文件？

你的代码核心错误是把csv.writer绑定到了原始的文件句柄fh，而非bz2包装后的文件对象bh。csv.writer需要操作的是经过bz2处理的文本模式流，而非直接操作磁盘文件句柄。

正确写法示例 #

用bz2.open()以文本模式打开文件，直接将这个对象传给csv.writer即可，无需手动打开底层文件：

import bz2
import csv

tfile = '/tmp/test.csv.bz'
row = ['bc22jtr', 118324, None, 'contran', None, 11.5, 9.23]

# 用bz2.open以文本模式打开，指定编码（推荐utf-8）
with bz2.open(tfile, mode='wt', compresslevel=9, encoding='utf-8') as bh:
    writer = csv.writer(bh)
    writer.writerow(row)
    # 可继续写入多行，比如writer.writerows(rows_list)

可行性说明 #

• bz2.open(tfile, 'wt')返回文本模式的类文件对象，会自动处理字符串到压缩字节流的转换，完全符合csv.writer对文本流的要求。
• csv.writer会自动处理字段的转义（换行符、逗号、特殊字符等），满足你的需求。
• 使用with语句可自动关闭文件，避免手动关闭时的错误，同时保证资源正确释放。

验证方法 #

写入后可通过以下代码验证文件是否正确：

import bz2
import csv

with bz2.open(tfile, mode='rt', encoding='utf-8') as bh:
    reader = csv.reader(bh)
    for row in reader:
        print(row)

输出应与写入的行一致，且所有特殊字符都被正确转义和还原。

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