Added way to turn off storing stdout and stderr in StdSim

Author: kmvanbrunt

cmd2/pyscript_bridge.py

@@ -25,9 +25,26 @@ class CommandResult(namedtuple_with_defaults('CommandResult', ['stdout', 'stderr
     Named tuple attributes
     ----------------------
     stdout: str - Output captured from stdout while this command is executing
-    stderr: str - Output captured from stderr while this command is executing. None if no error captured
+    stderr: str - Output captured from stderr while this command is executing. None if no error captured.
     data - Data returned by the command.
 
+    Any combination of these fields can be used when developing a scripting API for a given command.
+    By default stdout and stderr will be captured for you. If there is additional command specific data,
+    then write that to cmd2's _last_result member. That becomes the data member of this tuple.
+
+    In some cases, the data member may contain everything needed for a command and storing stdout
+    and stderr might just be a duplication of data that wastes memory. In that case, the StdSim can
+    be told not to store output with its set_store_output() method.
+
+    The code would look like this:
+        if isinstance(self.stdout, StdSim):
+            self.stdout.set_store_output(False)
+
+        if isinstance(sys.stderr, StdSim):
+            sys.stderr.set_store_output(False)
+
+    See StdSim class in utils.py for more information
+
     NOTE: Named tuples are immutable.  So the contents are there for access, not for modification.
     """
     def __bool__(self) -> bool:
@@ -67,25 +84,25 @@ def __call__(self, command: str, echo: Optional[bool] = None) -> CommandResult:
         if echo is None:
             echo = self.cmd_echo
 
-        copy_stdout = StdSim(sys.stdout, echo)
-        copy_stderr = StdSim(sys.stderr, echo)
-
+        # This will be used to capture _cmd2_app.stdout and sys.stdout
         copy_cmd_stdout = StdSim(self._cmd2_app.stdout, echo)
 
+        # This will be used to capture sys.stderr
+        copy_stderr = StdSim(sys.stderr, echo)
+
         self._cmd2_app._last_result = None
 
         try:
             self._cmd2_app.stdout = copy_cmd_stdout
-            with redirect_stdout(copy_stdout):
+            with redirect_stdout(copy_cmd_stdout):
                 with redirect_stderr(copy_stderr):
                     # Include a newline in case it's a multiline command
                     self._cmd2_app.onecmd_plus_hooks(command + '\n')
         finally:
             self._cmd2_app.stdout = copy_cmd_stdout.inner_stream
 
-        # if stderr is empty, set it to None
-        stderr = copy_stderr.getvalue() if copy_stderr.getvalue() else None
-
-        outbuf = copy_cmd_stdout.getvalue() if copy_cmd_stdout.getvalue() else copy_stdout.getvalue()
-        result = CommandResult(stdout=outbuf, stderr=stderr, data=self._cmd2_app._last_result)
+        # Save the output. If stderr is empty, set it to None.
+        result = CommandResult(stdout=copy_cmd_stdout.getvalue(),
+                               stderr=copy_stderr.getvalue() if copy_stderr.getvalue() else None,
+                               data=self._cmd2_app._last_result)
         return result

cmd2/utils.py

@@ -261,28 +261,10 @@ def natural_sort(list_to_sort: Iterable[str]) -> List[str]:
 
 
 class StdSim(object):
-    """Class to simulate behavior of sys.stdout or sys.stderr.
-
+    """
+    Class to simulate behavior of sys.stdout or sys.stderr.
     Stores contents in internal buffer and optionally echos to the inner stream it is simulating.
     """
-    class ByteBuf(object):
-        """Inner class which stores an actual bytes buffer and does the actual output if echo is enabled."""
-        def __init__(self, inner_stream, echo: bool = False,
-                     encoding: str = 'utf-8', errors: str = 'replace') -> None:
-            self.byte_buf = b''
-            self.inner_stream = inner_stream
-            self.echo = echo
-            self.encoding = encoding
-            self.errors = errors
-
-        def write(self, b: bytes) -> None:
-            """Add bytes to internal bytes buffer and if echo is True, echo contents to inner stream."""
-            if not isinstance(b, bytes):
-                raise TypeError('a bytes-like object is required, not {}'.format(type(b)))
-            self.byte_buf += b
-            if self.echo:
-                self.inner_stream.buffer.write(b)
-
     def __init__(self, inner_stream, echo: bool = False,
                  encoding: str = 'utf-8', errors: str = 'replace') -> None:
         """
@@ -292,17 +274,20 @@ def __init__(self, inner_stream, echo: bool = False,
         :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.echo = echo
         self.encoding = encoding
         self.errors = errors
+        self.__store_output = True
+        self.buffer = ByteBuf(self)
 
     def write(self, s: str) -> None:
         """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)))
-        self.buffer.byte_buf += s.encode(encoding=self.encoding, errors=self.errors)
+
+        if self.__store_output:
+            self.buffer.byte_buf += s.encode(encoding=self.encoding, errors=self.errors)
         if self.echo:
             self.inner_stream.write(s)
 
@@ -330,13 +315,42 @@ def clear(self) -> None:
         """Clear the internal contents"""
         self.buffer.byte_buf = b''
 
+    def get_store_output(self) -> bool:
+        return self.__store_output
+
+    def set_store_output(self, store_output: bool) -> None:
+        """
+        Set whether output should be saved in buffer.byte_buf
+        :param store_output: Store output if True, otherwise do not and clear the buffer
+        """
+        self.__store_output = self.buffer.store_output = store_output
+        self.clear()
+
     def __getattr__(self, item: str):
         if item in self.__dict__:
             return self.__dict__[item]
         else:
             return getattr(self.inner_stream, item)
 
 
+class ByteBuf(object):
+    """
+    Used by StdSim to write binary data and stores the actual bytes written
+    """
+    def __init__(self, std_sim_instance: StdSim) -> None:
+        self.byte_buf = b''
+        self.std_sim_instance = std_sim_instance
+
+    def write(self, b: bytes) -> None:
+        """Add bytes to internal bytes buffer and if echo is True, echo contents to inner stream."""
+        if not isinstance(b, bytes):
+            raise TypeError('a bytes-like object is required, not {}'.format(type(b)))
+        if self.std_sim_instance.get_store_output():
+            self.byte_buf += b
+        if self.std_sim_instance.echo:
+            self.std_sim_instance.inner_stream.buffer.write(b)
+
+
 def unquote_redirection_tokens(args: List[str]) -> None:
     """
     Unquote redirection tokens in a list of command-line arguments