Skip to content

Instantly share code, notes, and snippets.

@vadmium
Created September 3, 2016 13:46
Show Gist options
  • Save vadmium/d1b07d771fbf4347683c005c40991c02 to your computer and use it in GitHub Desktop.
Save vadmium/d1b07d771fbf4347683c005c40991c02 to your computer and use it in GitHub Desktop.
Mixin to add buffer attribute to text stream
from io import TextIOBase, RawIOBase, StringIO, BufferedRWPair
class _AsciiWrapper(RawIOBase):
'''Byte stream wrapping a TextIOBase stream via ASCII encoding'''
def __init__(self, text):
self._text = text
def close(self):
self._text.close()
def fileno(self):
return self._text.fileno()
def flush(self):
self._text.flush()
def isatty(self):
return self._text.isatty()
def readable(self):
return self._text.readable()
def writable(self):
return self._text.writable()
def truncate(self, size=None):
if size is not None:
return super().truncate(size)
return super(AsciiBufferMixin, self)._text.truncate()
def readinto(self, b):
with memoryview(b) as b, b.cast("B") as b:
result = self._text.read(len(b))
b[:len(result)] = result.encode("ascii")
return len(result)
def write(self, b):
return self._text.write(str(b, "ascii"))
class AsciiBufferMixin(TextIOBase):
'''Add "buffer" property that does ASCII encoding to a text stream
The "buffer" object handles newline translation the same as the text
stream, rather than bypassing the translation like TextIOWrapper.buffer
would.'''
encoding = "ascii"
errors = "strict"
def __init__(self, *pos, **kw):
super().__init__(*pos, **kw)
wrapper = _AsciiWrapper(self)
self.buffer = BufferedRWPair(wrapper, wrapper)
def detach(self):
return self.buffer
def close(self):
try:
self.buffer.flush()
finally:
super().close()
def seek(self, *args):
self.buffer.flush()
return super().seek(*args)
def tell(self):
self.buffer.flush()
return super().tell()
def trucate(self, *arg):
self.buffer.flush()
return super().truncate(*arg)
class Stream(AsciiBufferMixin, StringIO):
pass
stream = Stream(newline="\r\n")
stream.write("Text\n")
stream.buffer.write(b"Bytes\n")
stream.buffer.flush()
assert stream.getvalue() == "Text\r\n" "Bytes\r\n"
stream = Stream("Line 1\r\n" "Line 2\r\n", newline=None)
assert stream.readline() == "Line 1\n"
assert stream.buffer.readline() == b"Line 2\n"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment