Merge pull request #564 from python-cmd2/stdsim_binary

StdSim improvements for handling binary data

2 files changed, 36 insertions, 12 deletions

diff --git a/cmd2/utils.py b/cmd2/utils.py
index ddd43507..501733a9 100644
--- a/cmd2/utils.py
+++ b/cmd2/utils.py
@@ -268,35 +268,49 @@ class StdSim(object):
             if self.echo:
                 self.inner_stream.buffer.write(b)
 
-    def __init__(self, inner_stream, echo: bool = False) -> None:
+    def __init__(self, inner_stream, echo: bool = False,
+                 encoding: str='utf-8', errors: str='replace') -> None:
+        """
+        Initializer
+        :param inner_stream: the emulated stream
+        :param echo: if True, then all input will be echoed to inner_stream
+        :param encoding: codec for encoding/decoding strings (defaults to utf-8)
+        :param errors: how to handle encoding/decoding errors (defaults to replace)
+        """
         self.buffer = self.ByteBuf(inner_stream, echo)
         self.inner_stream = inner_stream
+        self.encoding = encoding
+        self.errors = errors
 
     def write(self, s: str) -> None:
-        """Add str to internal bytes buffer and if echo is True, echo contents to inner stream."""
+        """Add str to internal bytes buffer and if echo is True, echo contents to inner stream"""
         if not isinstance(s, str):
             raise TypeError('write() argument must be str, not {}'.format(type(s)))
-        b = s.encode()
+        b = s.encode(encoding=self.encoding, errors=self.errors)
         self.buffer.write(b)
 
     def getvalue(self) -> str:
-        """Get the internal contents as a str.
+        """Get the internal contents as a str"""
+        return self.buffer.byte_buf.decode(encoding=self.encoding, errors=self.errors)
 
-        :return string from the internal contents
-        """
-        return self.buffer.byte_buf.decode()
+    def getbytes(self) -> bytes:
+        """Get the internal contents as bytes"""
+        return self.buffer.byte_buf
 
     def read(self) -> str:
-        """Read from the internal contents as a str and then clear them out.
-
-        :return: string from the internal contents
-        """
+        """Read from the internal contents as a str and then clear them out"""
         result = self.getvalue()
         self.clear()
         return result
 
+    def readbytes(self) -> bytes:
+        """Read from the internal contents as bytes and then clear them out"""
+        result = self.getbytes()
+        self.clear()
+        return result
+
     def clear(self) -> None:
-        """Clear the internal contents."""
+        """Clear the internal contents"""
         self.buffer.byte_buf = b''
 
     def __getattr__(self, item: str):
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 53031567..43a05a9a 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -133,6 +133,7 @@ def test_stdsim_buffer_write_bytes(stdout_sim):
     b_str = b'Hello World'
     stdout_sim.buffer.write(b_str)
     assert stdout_sim.getvalue() == b_str.decode()
+    assert stdout_sim.getbytes() == b_str
 
 def test_stdsim_buffer_write_str(stdout_sim):
     my_str = 'Hello World'
@@ -148,6 +149,15 @@ def test_stdsim_read(stdout_sim):
     assert stdout_sim.read() == my_str
     assert stdout_sim.getvalue() == ''
 
+def test_stdsim_read_bytes(stdout_sim):
+    b_str = b'Hello World'
+    stdout_sim.buffer.write(b_str)
+    # getbytes() returns the value and leaves it unaffected internally
+    assert stdout_sim.getbytes() == b_str
+    # read_bytes() returns the value and then clears the internal buffer
+    assert stdout_sim.readbytes() == b_str
+    assert stdout_sim.getbytes() == b''
+
 def test_stdsim_clear(stdout_sim):
     my_str = 'Hello World'
     stdout_sim.write(my_str)