Checkpoint implementing PATCH according to Docker

util/registry/filelike.py

@@ -1,24 +1,103 @@
-class SocketReader(object):
-  def __init__(self, fp):
-    self._fp = fp
+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 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')
+
+    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)
+
+
+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, n=-1):
-    buf = self._fp.read(n)
-    if not buf:
-      return ''
+  def read(self, size=READ_UNTIL_END):
+    buf = super(SocketReader, self).read(size)
     for handler in self.handlers:
       handler(buf)
     return buf
 
-  def tell(self):
-    raise IOError('Stream is not seekable.')
-
 
 def wrap_with_handler(in_fp, handler):
   wrapper = SocketReader(in_fp)
   wrapper.add_handler(handler)
   return wrapper
+
+
+class FilelikeStreamConcat(BaseStreamFilelike):
+  def __init__(self, file_generator):
+    super(FilelikeStreamConcat, self).__init__(self)
+    self._file_generator = file_generator
+    self._current_file = file_generator.next()
+
+  def read(self, size=READ_UNTIL_END):
+    buf = self._current_file.read(size)
+    if buf:
+      self._cursor_position += len(buf)
+      return buf
+
+    # 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 ''
+    return self.read(size)
+
+
+class StreamSlice(BaseStreamFilelike):
+  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
+
+    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 - 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)