quay/util/registry/filelike.py

WHENCE_ABSOLUTE = 0
WHENCE_RELATIVE = 1
WHENCE_RELATIVE_END = 2

READ_UNTIL_END = -1


class BaseStreamFilelike(object):
  def __init__(self, fileobj):
    self._fileobj = fileobj
    self._cursor_position = 0

  def close(self):
    self._fileobj.close()

  def read(self, size=READ_UNTIL_END):
    buf = self._fileobj.read(size)
    self._cursor_position += len(buf)
    return buf

  def tell(self):
    return self._cursor_position

  def seek(self, index, whence=WHENCE_ABSOLUTE):
    num_bytes_to_ff = 0
    if whence == WHENCE_ABSOLUTE:
      if index < self._cursor_position:
        raise IOError('Cannot seek backwards')
      num_bytes_to_ff = index - self._cursor_position

    elif whence == WHENCE_RELATIVE:
      if index < 0:
        raise IOError('Cannnot seek backwards')
      num_bytes_to_ff = index

    elif whence == WHENCE_RELATIVE_END:
      raise IOError('Stream does not have a known end point')

    bytes_forward = num_bytes_to_ff
    while num_bytes_to_ff > 0:
      buf = self._fileobj.read(num_bytes_to_ff)
      if not buf:
        raise IOError('Seek past end of file')
      num_bytes_to_ff -= len(buf)

    self._cursor_position += bytes_forward
    return bytes_forward


class SocketReader(BaseStreamFilelike):
  def __init__(self, fileobj):
    super(SocketReader, self).__init__(fileobj)
    self.handlers = []

  def add_handler(self, handler):
    self.handlers.append(handler)

  def read(self, size=READ_UNTIL_END):
    buf = super(SocketReader, self).read(size)
    for handler in self.handlers:
      handler(buf)
    return buf


def wrap_with_handler(in_fp, handler):
  wrapper = SocketReader(in_fp)
  wrapper.add_handler(handler)
  return wrapper


class FilelikeStreamConcat(object):
  """ A file-like object which concats all the file-like objects in the specified generator into
      a single stream.
  """
  def __init__(self, file_generator):
    self._file_generator = file_generator
    self._current_file = file_generator.next()
    self._current_position = 0
    self._closed = False

  def tell(self):
    return self._current_position

  def close(self):
    self._closed = True

  def read(self, size=READ_UNTIL_END):
    buf = ''
    current_size = size

    while size == READ_UNTIL_END or len(buf) < size:
      current_buf = self._current_file.read(current_size)
      if current_buf:
        buf += current_buf
        self._current_position += len(current_buf)
        if size != READ_UNTIL_END:
          current_size -= len(current_buf)

      else:
        # That file was out of data, prime a new one
        self._current_file.close()
        try:
          self._current_file = self._file_generator.next()
        except StopIteration:
          return buf

    return buf


class StreamSlice(BaseStreamFilelike):
  """ A file-like object which returns a file-like object that represents a slice of the data in
      the specified file obj. All methods will act as if the slice is its own file.
  """

  def __init__(self, fileobj, start_offset=0, end_offset_exclusive=READ_UNTIL_END):
    super(StreamSlice, self).__init__(fileobj)
    self._end_offset_exclusive = end_offset_exclusive
    self._start_offset = start_offset

    if start_offset > 0:
      self.seek(start_offset)

  def read(self, size=READ_UNTIL_END):
    if self._end_offset_exclusive == READ_UNTIL_END:
      # We weren't asked to limit the end of the stream
      return super(StreamSlice, self).read(size)

    # Compute the max bytes to read until the end or until we reach the user requested max
    max_bytes_to_read = self._end_offset_exclusive - super(StreamSlice, self).tell()
    if size != READ_UNTIL_END:
      max_bytes_to_read = min(max_bytes_to_read, size)

    return super(StreamSlice, self).read(max_bytes_to_read)

  def _file_min(self, first, second):
    if first == READ_UNTIL_END:
      return second

    if second == READ_UNTIL_END:
      return first

    return min(first, second)

  def tell(self):
    return super(StreamSlice, self).tell() - self._start_offset

  def seek(self, index, whence=WHENCE_ABSOLUTE):
    index = self._file_min(self._end_offset_exclusive, index)
    super(StreamSlice, self).seek(index, whence)


class LimitingStream(StreamSlice):
  """ A file-like object which mimics the specified file stream being limited to the given number
      of bytes. All calls after that limit (if specified) will act as if the file has no additional
      data.
  """
  def __init__(self, fileobj, read_limit=READ_UNTIL_END):
    super(LimitingStream, self).__init__(fileobj, 0, read_limit)