diff --git a/storage/swift.py b/storage/swift.py index ddeae9105..cfb0be0d5 100644 --- a/storage/swift.py +++ b/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, diff --git a/test/test_util.py b/test/test_util.py index ae27c670a..8afa611c4 100644 --- a/test/test_util.py +++ b/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() + diff --git a/util/generatorfile.py b/util/generatorfile.py new file mode 100644 index 000000000..36544cc28 --- /dev/null +++ b/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 \ No newline at end of file