gh-129205: Experiment BytesIO._readfrom()

Include/internal/pycore_global_strings.h

@@ -405,6 +405,7 @@ struct _Py_global_strings{
  STRUCT_FOR_ID(entrypoint)
  STRUCT_FOR_ID(env)
  STRUCT_FOR_ID(errors)
+ STRUCT_FOR_ID(estimate)
  STRUCT_FOR_ID(event)
  STRUCT_FOR_ID(eventmask)
  STRUCT_FOR_ID(exc_type)

Lib/_compression.py

@@ -111,14 +111,10 @@ def read(self, size=-1):
  return data
 
  def readall(self):
- chunks = []
- # sys.maxsize means the max length of output buffer is unlimited,
- # so that the whole input buffer can be decompressed within one
- # .decompress() call.
- while data := self.read(sys.maxsize):
- chunks.append(data)
-
- return b"".join(chunks)
+ # FIXME(cmaloney): non blocking support?
+ bio = io.BytesIO()
+ bio.readfrom(self)
+ return bio.getvalue()
 
  # Rewind the file to the beginning of the data stream.
  def _rewind(self):

Lib/_pyio.py

@@ -926,6 +926,82 @@ def read1(self, size=-1):
  """
  return self.read(size)
 
+ def readfrom(self, file, /, *, estimate=None, limit=None):
+ """Efficiently read from the provided file and return True if hit end.
+
+ Returns True if and only if a read into a non-zero length buffer
+ returns 0 bytes. On most systems this indicates end of file / stream.
+ """
+ if self.closed:
+ raise ValueError("read from closed file")
+
+ # In order to detect end of file, need a read() of at least 1
+ # byte which returns size 0. Oversize the buffer by 1 byte so the
+ # I/O can be completed with two read() calls (one for all data, one
+ # for EOF) without needing to resize the buffer.
+ target_read = None
+ if estimate is not None:
+ target_read = int(estimate) + 1
+ else:
+ target_read = DEFAULT_BUFFER_SIZE
+
+ # Cap to limit
+ if limit is not None:
+ limit = int(limit)
+ if limit == 0: # Nothing to read.
+ return False
+ if limit < 0:
+ raise ValueError(f"limit must be larger than 0, got{limit}")
+
+ if limit is not None:
+ target_read = min(target_read, limit)
+
+ # Expand buffer to get target read in one read when possible.
+ if len(self._buffer) < target_read + self._pos:
+ self._buffer.resize(self._pos + target_read)
+
+ if isinstance(file, int): # File descriptor
+ read_fn = lambda: os.readinto(file, memoryview(self._buffer)[self._pos:])
+ elif file_readinto := getattr(file, "readinto", None):
+ read_fn = lambda: file_readinto(memoryview(self._buffer)[self._pos:])
+ elif file_read := getattr(file, "read", None):
+ def read_fn():
+ data = file_read(len(self._buffer) - self._pos)
+ self._buffer[self._pos:self._pos + len(data)] = data
+
+ found_eof = False
+ start_pos = self._pos
+ try:
+ while n := read_fn():
+ self._pos += n
+ # Expand buffer if needed.
+ if len(self._buffer) - self._pos <= 0:
+ bytes_read = self._pos - start_pos
+ target_read = _new_buffersize(bytes_read)
+
+ # Keep buffer size <= limit, so only need to check against
+ # limit when resizing.
+ if limit is not None:
+ remaining = limit - bytes_read
+ if remaining <= 0:
+ assert remaining == 0, "should never pass limit"
+ break
+ target_read = min(remaining, target_read)
+
+ self._buffer.resize(target_read + len(self._buffer))
+
+ else:
+ assert len(self._buffer) - self._pos >= 1, \
+ "os.readinto buffer size 0 will result in erroneous EOF / returns 0"
+ found_eof = True
+
+ except BlockingIOError:
+ pass
+
+ # Remove all excess bytes.
+ self._buffer.resize(self._pos)
+ return found_eof
+
  def write(self, b):
  if self.closed:
  raise ValueError("write to closed file")
@@ -1666,38 +1742,22 @@ def readall(self):
  """
  self._checkClosed()
  self._checkReadable()
- if self._stat_atopen is None or self._stat_atopen.st_size <= 0:
- bufsize = DEFAULT_BUFFER_SIZE
- else:
- # In order to detect end of file, need a read() of at least 1
- # byte which returns size 0. Oversize the buffer by 1 byte so the
- # I/O can be completed with two read() calls (one for all data, one
- # for EOF) without needing to resize the buffer.
- bufsize = self._stat_atopen.st_size + 1
-
- if self._stat_atopen.st_size > 65536:
+ estimate = None
+ if self._stat_atopen and self._stat_atopen.st_size >= 0:
+ estimate = self._stat_atopen.st_size
+ if estimate > 65536:
  try:
  pos = os.lseek(self._fd, 0, SEEK_CUR)
- if self._stat_atopen.st_size >= pos:
- bufsize = self._stat_atopen.st_size - pos + 1
+ estimate = estimate - pos if estimate > pos else 0
  except OSError:
  pass
 
- result = bytearray(bufsize)
- bytes_read = 0
- try:
- while n := os.readinto(self._fd, memoryview(result)[bytes_read:]):
- bytes_read += n
- if bytes_read >= len(result):
- result.resize(_new_buffersize(bytes_read))
- except BlockingIOError:
- if not bytes_read:
- return None
+ bio = BytesIO()
+ found_eof = bio.readfrom(self._fd, estimate=estimate)
+ result = bio.getvalue()
+ # No limit in readfrom, so not finding eof indicates blocked.
+ return result if result or found_eof else None
 
- assert len(result) - bytes_read >= 1, \
- "os.readinto buffer size 0 will result in erroneous EOF / returns 0"
- result.resize(bytes_read)
- return bytes(result)
 
  def readinto(self, buffer):
  """Same as RawIOBase.readinto()."""

Lib/subprocess.py

@@ -1921,12 +1921,9 @@ def _execute_child(self, args, executable, preexec_fn, close_fds,
 
  # Wait for exec to fail or succeed; possibly raising an
  # exception (limited in size)
- errpipe_data = bytearray()
- while True:
- part = os.read(errpipe_read, 50000)
- errpipe_data += part
- if not part or len(errpipe_data) > 50000:
- break
+ bio = io.BytesIO()
+ bio.readfrom(errpipe_read, estimate=0, limit=50_000)
+ errpipe_data = bio.getvalue()
  finally:
  # be sure the FD is closed no matter what
  os.close(errpipe_read)