diff options
| -rw-r--r-- | cmd2/pyscript_bridge.py | 37 | ||||
| -rw-r--r-- | cmd2/utils.py | 58 |
2 files changed, 63 insertions, 32 deletions
diff --git a/cmd2/pyscript_bridge.py b/cmd2/pyscript_bridge.py index 6c14ff1d..d2e52a30 100644 --- a/cmd2/pyscript_bridge.py +++ b/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 @@ class PyscriptBridge(object): 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 diff --git a/cmd2/utils.py b/cmd2/utils.py index 098ed41d..ca68bff7 100644 --- a/cmd2/utils.py +++ b/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 @@ class StdSim(object): :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,6 +315,17 @@ class StdSim(object): """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] @@ -337,6 +333,24 @@ class StdSim(object): 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 |