From 8c47a4a5cc9983b2775e2a5762e5eb48f47d6d08 Mon Sep 17 00:00:00 2001
From: Kevin Van Brunt <kmvanbrunt@gmail.com>
Date: Wed, 3 Oct 2018 11:15:30 -0400
Subject: Fixed error where binary data was crashing getvalue. Added accessor
 methods for retrieving StdSim contents as bytes.

---
 cmd2/utils.py | 40 +++++++++++++++++++++++++++-------------
 1 file changed, 27 insertions(+), 13 deletions(-)

(limited to 'cmd2')

diff --git a/cmd2/utils.py b/cmd2/utils.py
index ddd43507..9b043178 100644
--- a/cmd2/utils.py
+++ b/cmd2/utils.py
@@ -272,31 +272,45 @@ class StdSim(object):
         self.buffer = self.ByteBuf(inner_stream, echo)
         self.inner_stream = inner_stream
 
-    def write(self, s: str) -> None:
-        """Add str to internal bytes buffer and if echo is True, echo contents to inner stream."""
+    def write(self, s: str, encoding: str = 'utf-8') -> None:
+        """
+        Add str to internal bytes buffer and if echo is True, echo contents to inner stream.
+        :param: s: string to write
+        :param encoding: how to encode the string (defaults to utf-8)
+        """
         if not isinstance(s, str):
             raise TypeError('write() argument must be str, not {}'.format(type(s)))
-        b = s.encode()
+        b = s.encode(encoding=encoding)
         self.buffer.write(b)
 
-    def getvalue(self) -> str:
-        """Get the internal contents as a str.
-
-        :return string from the internal contents
+    def getvalue(self, encoding: str = 'utf-8') -> str:
+        """
+        Get the internal contents as a str
+        :param encoding: how to decode the bytes (defaults to utf-8)
         """
-        return self.buffer.byte_buf.decode()
+        return self.buffer.byte_buf.decode(encoding=encoding, errors='replace')
 
-    def read(self) -> str:
-        """Read from the internal contents as a str and then clear them out.
+    def getbytes(self) -> bytes:
+        """Get the internal contents as bytes"""
+        return self.buffer.byte_buf
 
-        :return: string from the internal contents
+    def read(self, encoding: str = 'utf-8') -> str:
         """
-        result = self.getvalue()
+        Read from the internal contents as a str and then clear them out
+        :param encoding: how to decode the bytes (defaults to utf-8)
+        """
+        result = self.getvalue(encoding)
+        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):
-- 
cgit v1.2.1