Merge pull request #214 from coreos-inc/swiftstream

Implement `stream_read_file` for the Swift storage engine

storage/swift.py

@@ -1,6 +1,7 @@
 """ Swift storage driver. Based on: github.com/bacongobbler/docker-registry-driver-swift/ """
 from swiftclient.client import Connection, ClientException
 from storage.basestorage import BaseStorage
+from util.generatorfile import GeneratorFile
 
 from random import SystemRandom
 import string
@@ -8,6 +9,7 @@ import logging
 
 logger = logging.getLogger(__name__)
 
+
 class SwiftStorage(BaseStorage):
   def __init__(self, swift_container, storage_path, auth_url, swift_user,
                swift_password, auth_version=None, os_options=None, ca_cert_path=None):
@@ -143,7 +145,7 @@ class SwiftStorage(BaseStorage):
       yield data
 
   def stream_read_file(self, path):
-    raise NotImplementedError
+    return GeneratorFile(self.stream_read(path))
 
   def stream_write(self, path, fp, content_type=None, content_encoding=None):
     self._put_object(path, fp, self.buffer_size, content_type=content_type,

test/test_util.py

@@ -3,6 +3,49 @@ import unittest
 from itertools import islice
 
 from util.validation import generate_valid_usernames
+from util.generatorfile import GeneratorFile
+
+class TestGeneratorFile(unittest.TestCase):
+  def sample_generator(self):
+    yield 'this'
+    yield 'is'
+    yield 'a'
+    yield 'test'
+
+  def test_basic_generator(self):
+    with GeneratorFile(self.sample_generator()) as f:
+      self.assertEquals("thisisatest", f.read())
+
+  def test_same_lengths(self):
+    with GeneratorFile(self.sample_generator()) as f:
+      self.assertEquals("this", f.read(4))
+      self.assertEquals("is", f.read(2))
+      self.assertEquals("a", f.read(1))
+      self.assertEquals("test", f.read(4))
+
+  def test_indexed_lengths(self):
+    with GeneratorFile(self.sample_generator()) as f:
+      self.assertEquals("thisis", f.read(6))
+      self.assertEquals("atest", f.read(5))
+
+  def test_misindexed_lengths(self):
+    with GeneratorFile(self.sample_generator()) as f:
+      self.assertEquals("thisis", f.read(6))
+      self.assertEquals("ate", f.read(3))
+      self.assertEquals("st", f.read(2))
+      self.assertEquals("", f.read(2))
+
+  def test_misindexed_lengths_2(self):
+    with GeneratorFile(self.sample_generator()) as f:
+      self.assertEquals("thisisat", f.read(8))
+      self.assertEquals("e", f.read(1))
+      self.assertEquals("st", f.read(2))
+      self.assertEquals("", f.read(2))
+
+  def test_overly_long(self):
+    with GeneratorFile(self.sample_generator()) as f:
+      self.assertEquals("thisisatest", f.read(60))
+
 
 class TestUsernameGenerator(unittest.TestCase):
   def assert_generated_output(self, input_username, expected_output):
@@ -48,3 +91,8 @@ class TestUsernameGenerator(unittest.TestCase):
     self.assertEquals('a__0', generated_output[1])
     self.assertEquals('a__1', generated_output[2])
     self.assertEquals('a__2', generated_output[3])
+
+
+if __name__ == '__main__':
+  unittest.main()
+

util/generatorfile.py

@@ -0,0 +1,76 @@
+def _complain_ifclosed(closed):
+  if closed:
+    raise ValueError, "I/O operation on closed file"
+
+class GeneratorFile(object):
+  """ File-like object which wraps a Python generator to produce the file contents.
+      Modeled on StringIO and comments on the file-like interface copied from there.
+  """
+  def __init__(self, generator):
+    self._generator = generator
+    self._closed = False
+    self._buf = ''
+
+  def __iter__(self):
+    return self
+
+  def next(self):
+    """A file object is its own iterator, for example iter(f) returns f
+    (unless f is closed). When a file is used as an iterator, typically
+    in a for loop (for example, for line in f: print line), the next()
+    method is called repeatedly. This method returns the next input line,
+    or raises StopIteration when EOF is hit.
+    """
+    _complain_ifclosed(self._closed)
+    r = self.read()
+    if not r:
+      raise StopIteration
+    return r
+
+  def readline(self):
+    buf = []
+    while True:
+      c = self.read(size=1)
+      buf.append(c)
+      if c == '\n' or c == '':
+        return ''.join(buf)
+
+  def flush(self):
+    _complain_ifclosed(self._closed)
+
+  def read(self, size=-1):
+    """Read at most size bytes from the file
+    (less if the read hits EOF before obtaining size bytes).
+
+    If the size argument is negative or omitted, read all data until EOF
+    is reached. The bytes are returned as a string object. An empty
+    string is returned when EOF is encountered immediately.
+    """
+    _complain_ifclosed(self._closed)
+    buf = self._buf
+    while size < 0 or len(buf) < size:
+      try:
+        buf = buf + self._generator.next()
+      except StopIteration:
+        break
+
+    returned = ''
+    if size >= 1:
+      self._buf = buf[size:]
+      returned = buf[:size]
+    else:
+      self._buf = ''
+      returned = buf
+
+    return returned
+
+
+  def close(self):
+    self._closed = True
+    del self._buf
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, type, value, traceback):
+    self._closed = True