class IteratorAsBinaryFile(object): """Custom bundle iterator for streaming. <requests_toolbelt.streaming_iterator> """ def __init__(self, size, iterator, md5, encoding="utf-8"): #: The expected size of the upload self.size = int(size) self.md5 = md5 if self.size < 0: raise ValueError( "The size of the upload must be a positive integer") #: Attribute that requests will check to determine the length of the #: body. See bug #80 for more details self.len = self.size #: The iterator used to generate the upload data self.iterator = iterator #: Encoding the iterator is using self.encoding = encoding # The buffer we use to provide the correct number of bytes requested # during a read self._buffer = CustomBytesIO() def _get_bytes(self): try: return encode_with(next(self.iterator), self.encoding) except StopIteration: return b"" def _load_bytes(self, size): self._buffer.smart_truncate() amount_to_load = size - super_len(self._buffer) bytes_to_append = True while amount_to_load > 0 and bytes_to_append: bytes_to_append = self._get_bytes() amount_to_load -= self._buffer.append(bytes_to_append) def read(self, size=-1): size = int(size) if size == -1: return b"".join(self.iterator) self._load_bytes(size) s = self._buffer.read(size) if not s: self.len = 0 if size < 0: self.len = 0 self.md5.update(s) return s
class TestCustomBytesIO(unittest.TestCase): def setUp(self): self.instance = CustomBytesIO() def test_writable(self): assert hasattr(self.instance, 'write') assert self.instance.write(b'example') == 7 def test_readable(self): assert hasattr(self.instance, 'read') assert self.instance.read() == b'' assert self.instance.read(10) == b'' def test_can_read_after_writing_to(self): self.instance.write(b'example text') self.instance.read() == b'example text' def test_can_read_some_after_writing_to(self): self.instance.write(b'example text') self.instance.read(6) == b'exampl' def test_can_get_length(self): self.instance.write(b'example') self.instance.seek(0, 0) assert len(self.instance) == 7 def test_truncates_intelligently(self): self.instance.write(b'abcdefghijklmnopqrstuvwxyzabcd') # 30 bytes assert self.instance.tell() == 30 self.instance.seek(-10, 2) self.instance.smart_truncate() assert len(self.instance) == 10 assert self.instance.read() == b'uvwxyzabcd' assert self.instance.tell() == 10 def test_accepts_encoded_strings_with_unicode(self): """Accepts a string with encoded unicode characters.""" s = b'this is a unicode string: \xc3\xa9 \xc3\xa1 \xc7\xab \xc3\xb3' self.instance = CustomBytesIO(s) assert self.instance.read() == s
class TestCustomBytesIO(unittest.TestCase): def setUp(self): self.instance = CustomBytesIO() def test_writable(self): assert hasattr(self.instance, 'write') assert self.instance.write(b'example') == 7 def test_readable(self): assert hasattr(self.instance, 'read') assert self.instance.read() == b'' assert self.instance.read(10) == b'' def test_can_read_after_writing_to(self): self.instance.write(b'example text') self.instance.read() == b'example text' def test_can_read_some_after_writing_to(self): self.instance.write(b'example text') self.instance.read(6) == b'exampl' def test_can_get_length(self): self.instance.write(b'example') self.instance.seek(0, 0) assert self.instance.len == 7 def test_truncates_intelligently(self): self.instance.write(b'abcdefghijklmnopqrstuvwxyzabcd') # 30 bytes assert self.instance.tell() == 30 self.instance.seek(-10, 2) self.instance.smart_truncate() assert self.instance.len == 10 assert self.instance.read() == b'uvwxyzabcd' assert self.instance.tell() == 10 def test_accepts_encoded_strings_with_unicode(self): """Accepts a string with encoded unicode characters.""" s = b'this is a unicode string: \xc3\xa9 \xc3\xa1 \xc7\xab \xc3\xb3' self.instance = CustomBytesIO(s) assert self.instance.read() == s