Merge pull request #534 from python-cmd2/copy_stream

Fixed several hack classes built to simulate file descriptors

Author: tleonhardt

cmd2/pyscript_bridge.py

@@ -12,15 +12,15 @@
 import sys
 from typing import List, Callable
 
+from .argparse_completer import _RangeAction
+from .utils import namedtuple_with_defaults, StdSim
+
 # Python 3.4 require contextlib2 for temporarily redirecting stderr and stdout
 if sys.version_info < (3, 5):
     from contextlib2 import redirect_stdout, redirect_stderr
 else:
     from contextlib import redirect_stdout, redirect_stderr
 
-from .argparse_completer import _RangeAction
-from .utils import namedtuple_with_defaults
-
 
 class CommandResult(namedtuple_with_defaults('CommandResult', ['stdout', 'stderr', 'data'])):
     """Encapsulates the results from a command.
@@ -38,37 +38,12 @@ def __bool__(self):
         return not self.stderr and self.data is not None
 
 
-class CopyStream(object):
-    """Copies all data written to a stream"""
-    def __init__(self, inner_stream, echo: bool = False) -> None:
-        self.buffer = ''
-        self.inner_stream = inner_stream
-        self.echo = echo
-
-    def write(self, s):
-        self.buffer += s
-        if self.echo:
-            self.inner_stream.write(s)
-
-    def read(self):
-        raise NotImplementedError
-
-    def clear(self):
-        self.buffer = ''
-
-    def __getattr__(self, item: str):
-        if item in self.__dict__:
-            return self.__dict__[item]
-        else:
-            return getattr(self.inner_stream, item)
-
-
 def _exec_cmd(cmd2_app, func: Callable, echo: bool):
     """Helper to encapsulate executing a command and capturing the results"""
-    copy_stdout = CopyStream(sys.stdout, echo)
-    copy_stderr = CopyStream(sys.stderr, echo)
+    copy_stdout = StdSim(sys.stdout, echo)
+    copy_stderr = StdSim(sys.stderr, echo)
 
-    copy_cmd_stdout = CopyStream(cmd2_app.stdout, echo)
+    copy_cmd_stdout = StdSim(cmd2_app.stdout, echo)
 
     cmd2_app._last_result = None
 
@@ -81,9 +56,9 @@ def _exec_cmd(cmd2_app, func: Callable, echo: bool):
         cmd2_app.stdout = copy_cmd_stdout.inner_stream
 
     # if stderr is empty, set it to None
-    stderr = copy_stderr.buffer if copy_stderr.buffer else None
+    stderr = copy_stderr.getvalue() if copy_stderr.getvalue() else None
 
-    outbuf = copy_cmd_stdout.buffer if copy_cmd_stdout.buffer else copy_stdout.buffer
+    outbuf = copy_cmd_stdout.getvalue() if copy_cmd_stdout.getvalue() else copy_stdout.getvalue()
     result = CommandResult(stdout=outbuf, stderr=stderr, data=cmd2_app._last_result)
     return result
 

cmd2/transcript.py

@@ -44,7 +44,7 @@ def setUp(self):
 
             # Trap stdout
             self._orig_stdout = self.cmdapp.stdout
-            self.cmdapp.stdout = OutputTrap()
+            self.cmdapp.stdout = utils.StdSim(self.cmdapp.stdout)
 
     def runTest(self):  # was testall
         if self.cmdapp:
@@ -203,24 +203,3 @@ def tearDown(self):
         if self.cmdapp:
             # Restore stdout
             self.cmdapp.stdout = self._orig_stdout
-
-class OutputTrap(object):
-    """Instantiate an OutputTrap to divert/capture ALL stdout output.
-    For use in transcript testing.
-    """
-
-    def __init__(self):
-        self.contents = ''
-
-    def write(self, txt: str):
-        """Add text to the internal contents."""
-        self.contents += txt
-
-    def read(self) -> str:
-        """Read from the internal contents and then clear them out.
-
-        :return: str - text from the internal contents
-        """
-        result = self.contents
-        self.contents = ''
-        return result

cmd2/utils.py

@@ -246,3 +246,61 @@ def natural_sort(list_to_sort: List[str]) -> List[str]:
     :return: the list sorted naturally
     """
     return sorted(list_to_sort, key=natural_keys)
+
+
+class StdSim(object):
+    """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) -> None:
+            self.byte_buf = b''
+            self.inner_stream = inner_stream
+            self.echo = echo
+
+        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) -> None:
+        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."""
+        if not isinstance(s, str):
+            raise TypeError('write() argument must be str, not {}'.format(type(s)))
+        b = s.encode()
+        self.buffer.write(b)
+
+    def getvalue(self) -> str:
+        """Get the internal contents as a str.
+
+        :return string from the internal contents
+        """
+        return self.buffer.byte_buf.decode()
+
+    def read(self) -> str:
+        """Read from the internal contents as a str and then clear them out.
+
+        :return: string from the internal contents
+        """
+        result = self.getvalue()
+        self.clear()
+        return result
+
+    def clear(self) -> None:
+        """Clear the internal contents."""
+        self.buffer.byte_buf = b''
+
+    def __getattr__(self, item: str):
+        if item in self.__dict__:
+            return self.__dict__[item]
+        else:
+            return getattr(self.inner_stream, item)

tests/conftest.py

@@ -12,6 +12,7 @@
 from pytest import fixture
 
 import cmd2
+from cmd2.utils import StdSim
 
 # Prefer statically linked gnureadline if available (for macOS compatibility due to issues with libedit)
 try:
@@ -115,21 +116,6 @@
 """.format(color_str)
 
 
-class StdOut(object):
-    """ Toy class for replacing self.stdout in cmd2.Cmd instances for unit testing. """
-    def __init__(self):
-        self.buffer = ''
-
-    def write(self, s):
-        self.buffer += s
-
-    def read(self):
-        raise NotImplementedError
-
-    def clear(self):
-        self.buffer = ''
-
-
 def normalize(block):
     """ Normalize a block of text to perform comparison.
 
@@ -142,18 +128,18 @@ def normalize(block):
 
 
 def run_cmd(app, cmd):
-    """ Clear StdOut buffer, run the command, extract the buffer contents, """
+    """ Clear StdSim buffer, run the command, extract the buffer contents, """
     app.stdout.clear()
     app.onecmd_plus_hooks(cmd)
-    out = app.stdout.buffer
+    out = app.stdout.getvalue()
     app.stdout.clear()
     return normalize(out)
 
 
 @fixture
 def base_app():
     c = cmd2.Cmd()
-    c.stdout = StdOut()
+    c.stdout = StdSim(c.stdout)
     return c
 
 

tests/test_argparse.py

@@ -6,9 +6,9 @@
 import pytest
 
 import cmd2
-from unittest import mock
+from cmd2.utils import StdSim
 
-from .conftest import run_cmd, StdOut
+from .conftest import run_cmd
 
 # Prefer statically linked gnureadline if available (for macOS compatibility due to issues with libedit)
 try:
@@ -115,7 +115,7 @@ def do_talk(self, args, extra):
 @pytest.fixture
 def argparse_app():
     app = ArgparseApp()
-    app.stdout = StdOut()
+    app.stdout = StdSim(app.stdout)
     return app
 
 
@@ -217,12 +217,12 @@ def do_base(self, args):
             func(self, args)
         else:
             # No subcommand was provided, so call help
-            self.do_help('base')
+            self.do_help(['base'])
 
 @pytest.fixture
 def subcommand_app():
     app = SubcommandApp()
-    app.stdout = StdOut()
+    app.stdout = StdSim(app.stdout)
     return app
 
 

tests/test_autocompletion.py

@@ -1,20 +1,22 @@
+# coding=utf-8
 """
 Unit/functional testing for argparse completer in cmd2
 
 Copyright 2018 Eric Lin <[email protected]>
 Released under MIT license, see LICENSE file
 """
 import pytest
-from .conftest import run_cmd, normalize, StdOut, complete_tester
+
+from cmd2.utils import StdSim
+from .conftest import run_cmd, normalize, complete_tester
 
 from examples.tab_autocompletion import TabCompleteExample
 
 @pytest.fixture
 def cmd2_app():
-    c = TabCompleteExample()
-    c.stdout = StdOut()
-
-    return c
+    app = TabCompleteExample()
+    app.stdout = StdSim(app.stdout)
+    return app
 
 
 SUGGEST_HELP = '''Usage: suggest -t {movie, show} [-h] [-d DURATION{1..2}]