为 boto3 的 StreamingBody 模拟类似文件的 gzipped csv

Question

我真正的 S3 助手会执行以下操作：

def read_gzipped_csv_from_s3(self, key):
    return self.bucket.Object(key).get()

obj = S3Helper().read_gzipped_csv_from_s3(key)
df = pd.read_csv(obj['Body'], compression='gzip')

我需要模拟

read_gzipped_csv_from_s3()

方法进行单元测试。问题是响应应该是一个 gzipped CSV，我必须从字符串构造它，因为当测试在 Gitlab 的管道中运行时我无法存储任何内容。

所以我有一些 csv 作为字符串：

CSV_DATA = """
name,value,control
ABC,1.0,1
DEF,2.0,0
GHI,3.0,-1
"""

然后我有一些使用常规 CSV 文件来模拟 botocore.response.StreamingBody 的示例代码：

body_encoded = open('accounts.csv').read().encode()
mock_stream = StreamingBody(io.BytesIO(body_encoded), len(body_encoded))

但我不知道如何在内存中创建压缩的 CSV：我在某处找到了开头：

import gzip

buffer = BytesIO()
with gzip.GzipFile(fileobj=buffer, mode='wb') as compressed:
    with TextIOWrapper(compressed, encoding='utf-8') as wrapper:
        <can't figure out what's here>

非常感谢您的帮助。

尝试了 SO 中的大量其他片段并对其进行了修改，但没有成功。我期望的是：gzipped CSV 文件状对象传递给 StreamingBody

Answer 1

您可以使用

.write()

将数据写入

BytesIO

对象。您还需要

.seek()

将文件位置重置到开头才能读取。

import gzip
from io import BytesIO, TextIOWrapper

buffer = BytesIO()
with gzip.GzipFile(fileobj=buffer, mode='wb') as compressed:
    with TextIOWrapper(compressed, encoding='utf-8') as wrapper:
        wrapper.write(CSV_DATA)
buffer.seek(0)
df = pd.read_csv(buffer, compression='gzip')