gh-89550: Buffer GzipFile.write to reduce execution time by ~15%

Lib/gzip.py

@@ -22,6 +22,7 @@
 _COMPRESS_LEVEL_BEST = 9
 
 READ_BUFFER_SIZE = 128 * 1024
+_WRITE_BUFFER_SIZE = 4 * io.DEFAULT_BUFFER_SIZE
 
 
 def open(filename, mode="rb", compresslevel=_COMPRESS_LEVEL_BEST,
@@ -120,6 +121,21 @@ class BadGzipFile(OSError):
  """Exception raised in some cases for invalid gzip files."""
 
 
+class _WriteBufferStream(io.RawIOBase):
+ """Minimal object to pass WriteBuffer flushes into GzipFile"""
+ def __init__(self, gzip_file):
+ self.gzip_file = gzip_file
+
+ def write(self, data):
+ return self.gzip_file._write_raw(data)
+
+ def seekable(self):
+ return False
+
+ def writable(self):
+ return True
+
+
 class GzipFile(_compression.BaseStream):
  """The GzipFile class simulates most of the methods of a file object with
  the exception of the truncate() method.
@@ -184,6 +200,7 @@ def __init__(self, filename=None, mode=None,
  if mode is None:
  mode = getattr(fileobj, 'mode', 'rb')
 
+
  if mode.startswith('r'):
  self.mode = READ
  raw = _GzipReader(fileobj)
@@ -206,6 +223,9 @@ def __init__(self, filename=None, mode=None,
  zlib.DEF_MEM_LEVEL,
  0)
  self._write_mtime = mtime
+ self._buffer_size = _WRITE_BUFFER_SIZE
+ self._buffer = io.BufferedWriter(_WriteBufferStream(self),
+ buffer_size=self._buffer_size)
  else:
  raise ValueError("Invalid mode: {!r}".format(mode))
 
@@ -231,6 +251,11 @@ def _init_write(self, filename):
  self.bufsize = 0
  self.offset = 0 # Current file offset for seek(), tell(), etc
 
+ def tell(self):
+ self._check_not_closed()
+ self._buffer.flush()
+ return super().tell()
+
  def _write_gzip_header(self, compresslevel):
  self.fileobj.write(b'\037\213') # magic header
  self.fileobj.write(b'\010') # compression method
@@ -272,6 +297,10 @@ def write(self,data):
  if self.fileobj is None:
  raise ValueError("write() on closed GzipFile object")
 
+ return self._buffer.write(data)
+
+ def _write_raw(self, data):
+ # Called by our self._buffer underlying WriteBufferStream.
  if isinstance(data, (bytes, bytearray)):
  length = len(data)
  else:
@@ -322,16 +351,17 @@ def close(self):
  fileobj = self.fileobj
  if fileobj is None:
  return
- self.fileobj = None
  try:
  if self.mode == WRITE:
+ self._buffer.flush()
  fileobj.write(self.compress.flush())
  write32u(fileobj, self.crc)
  # self.size may exceed 2 GiB, or even 4 GiB
  write32u(fileobj, self.size & 0xffffffff)
  elif self.mode == READ:
  self._buffer.close()
  finally:
+ self.fileobj = None
  myfileobj = self.myfileobj
  if myfileobj:
  self.myfileobj = None
@@ -341,7 +371,7 @@ def flush(self,zlib_mode=zlib.Z_SYNC_FLUSH):
  self._check_not_closed()
  if self.mode == WRITE:
  # Ensure the compressor's buffer is flushed
- self.fileobj.write(self.compress.flush(zlib_mode))
+ self._buffer.flush()
  self.fileobj.flush()
 
  def fileno(self):
@@ -378,10 +408,10 @@ def seek(self, offset, whence=io.SEEK_SET):
  if offset < self.offset:
  raise OSError('Negative seek in write mode')
  count = offset - self.offset
- chunk = b'\0' * 1024
- for i in range(count // 1024):
+ chunk = b'\0' * self._buffer_size
+ for i in range(count // self._buffer_size):
  self.write(chunk)
- self.write(b'\0' * (count % 1024))
+ self.write(b'\0' * (count % self._buffer_size))
  elif self.mode == READ:
  self._check_not_closed()
  return self._buffer.seek(offset, whence)

Misc/NEWS.d/next/Library/2023-01-22-14-53-12.gh-issue-89550.c1U23f.rst

@@ -0,0 +1,2 @@
+Decrease execution time of some :mod:`gzip` file writes by 15% by
+adding more appropriate buffering.