From 6d1e39dbdca770a30a9be7df6e7a027fbe1d54a9 Mon Sep 17 00:00:00 2001 From: Tom Forbes Date: Sun, 20 Oct 2024 11:33:59 +0100 Subject: [PATCH] Buffer correctly --- src/pypi_data/cli.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/pypi_data/cli.py b/src/pypi_data/cli.py index 43dc032f..c0a8981e 100644 --- a/src/pypi_data/cli.py +++ b/src/pypi_data/cli.py @@ -1,6 +1,7 @@ import asyncio import contextlib import gzip +import io from collections import defaultdict from pathlib import Path from typing import ( @@ -56,11 +57,13 @@ def github_client(github_token) -> Github: @contextlib.contextmanager def open_path(path: Path, mode: Literal["wb", "rb"]) -> Generator[BinaryIO, None, None]: buffer_size = 1024 * 1024 * 50 - with path.open(mode, buffering=buffer_size) as fd: - if path.suffix == ".gz": - with gzip.open(fd, mode) as gzip_fd: - yield gzip_fd - else: + + if path.suffix == ".gz": + with gzip.open(path, mode) as gzip_fd: + with io.BufferedWriter(gzip_fd, buffer_size=buffer_size) as buffered_fd: + yield buffered_fd + else: + with path.open(mode, buffering=buffer_size) as fd: yield fd log.info(