Add allow_clipboard for #1225

3 files changed, 67 insertions, 51 deletions

diff --git a/cmd2/clipboard.py b/cmd2/clipboard.py
index fa1c23a7..9adea7ac 100644
--- a/cmd2/clipboard.py
+++ b/cmd2/clipboard.py
@@ -2,37 +2,16 @@
 """
 This module provides basic ability to copy from and paste to the clipboard/pastebuffer.
 """
-from typing import (
-    cast,
-)
-
+import typing
 import pyperclip  # type: ignore[import]
 
-# noinspection PyProtectedMember
-
-# Can we access the clipboard?  Should always be true on Windows and Mac, but only sometimes on Linux
-# noinspection PyBroadException
-try:
-    # Try getting the contents of the clipboard
-    _ = pyperclip.paste()
-
-# pyperclip raises at least the following types of exceptions. To be safe, just catch all Exceptions.
-#   FileNotFoundError on Windows Subsystem for Linux (WSL) when Windows paths are removed from $PATH
-#   ValueError for headless Linux systems without Gtk installed
-#   AssertionError can be raised by paste_klipper().
-#   PyperclipException for pyperclip-specific exceptions
-except Exception:
-    can_clip = False
-else:
-    can_clip = True
-
 
 def get_paste_buffer() -> str:
     """Get the contents of the clipboard / paste buffer.
 
     :return: contents of the clipboard
     """
-    pb_str = cast(str, pyperclip.paste())
+    pb_str = typing.cast(str, pyperclip.paste())
     return pb_str
 
 
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 6cd8b950..4f428d1b 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -37,6 +37,8 @@ import os
 import pydoc
 import re
 import sys
+import tempfile
+import typing
 import threading
 from code import (
     InteractiveConsole,
@@ -83,7 +85,6 @@ from .argparse_custom import (
     CompletionItem,
 )
 from .clipboard import (
-    can_clip,
     get_paste_buffer,
     write_to_paste_buffer,
 )
@@ -236,6 +237,7 @@ class Cmd(cmd.Cmd):
         shortcuts: Optional[Dict[str, str]] = None,
         command_sets: Optional[Iterable[CommandSet]] = None,
         auto_load_commands: bool = True,
+        allow_clipboard: bool = True,
     ) -> None:
         """An easy but powerful framework for writing line-oriented command
         interpreters. Extends Python's cmd package.
@@ -283,6 +285,7 @@ class Cmd(cmd.Cmd):
                                    that are currently loaded by Python and automatically
                                    instantiate and register all commands. If False, CommandSets
                                    must be manually installed with `register_command_set`.
+        :param allow_clipboard: If False, cmd2 will disable clipboard interactions
         """
         # Check if py or ipy need to be disabled in this instance
         if not include_py:
@@ -436,8 +439,8 @@ class Cmd(cmd.Cmd):
             self.pager = 'less -RXF'
             self.pager_chop = 'less -SRXF'
 
-        # This boolean flag determines whether or not the cmd2 application can interact with the clipboard
-        self._can_clip = can_clip
+        # This boolean flag stores whether cmd2 will allow clipboard related features
+        self.allow_clipboard = allow_clipboard
 
         # This determines the value returned by cmdloop() when exiting the application
         self.exit_code = 0
@@ -2734,13 +2737,9 @@ class Cmd(cmd.Cmd):
                 sys.stdout = self.stdout = new_stdout
 
         elif statement.output:
-            import tempfile
 
-            if (not statement.output_to) and (not self._can_clip):
-                raise RedirectionError("Cannot redirect to paste buffer; missing 'pyperclip' and/or pyperclip dependencies")
-
-            # Redirecting to a file
-            elif statement.output_to:
+            if statement.output_to:
+                # redirecting to a file
                 # statement.output can only contain REDIRECTION_APPEND or REDIRECTION_OUTPUT
                 mode = 'a' if statement.output == constants.REDIRECTION_APPEND else 'w'
                 try:
@@ -2752,14 +2751,26 @@ class Cmd(cmd.Cmd):
                 redir_saved_state.redirecting = True
                 sys.stdout = self.stdout = new_stdout
 
-            # Redirecting to a paste buffer
             else:
+                # Redirecting to a paste buffer
+                # we are going to direct output to a temporary file, then read it back in and
+                # put it in the paste buffer later
+                if not self.allow_clipboard:
+                    raise RedirectionError("Clipboard access not allowed")
+
+                # attempt to get the paste buffer, this forces pyperclip to go figure
+                # out if it can actually interact with the paste buffer, and will throw exceptions
+                # if it's not gonna work. That way we throw the exception before we go
+                # run the command and queue up all the output. if this is going to fail,
+                # no point opening up the temporary file
+                current_paste_buffer = get_paste_buffer()
+                # create a temporary file to store output
                 new_stdout = cast(TextIO, tempfile.TemporaryFile(mode="w+"))
                 redir_saved_state.redirecting = True
                 sys.stdout = self.stdout = new_stdout
 
                 if statement.output == constants.REDIRECTION_APPEND:
-                    self.stdout.write(get_paste_buffer())
+                    self.stdout.write(current_paste_buffer)
                     self.stdout.flush()
 
         # These are updated regardless of whether the command redirected
@@ -4551,8 +4562,6 @@ class Cmd(cmd.Cmd):
                 self.last_result = True
                 return stop
         elif args.edit:
-            import tempfile
-
             fd, fname = tempfile.mkstemp(suffix='.txt', text=True)
             fobj: TextIO
             with os.fdopen(fd, 'w') as fobj:
diff --git a/tests/test_cmd2.py b/tests/test_cmd2.py
index 5bcceb34..859f035a 100755
--- a/tests/test_cmd2.py
+++ b/tests/test_cmd2.py
@@ -729,8 +729,20 @@ def test_pipe_to_shell_error(base_app):
     assert not out
     assert "Pipe process exited with code" in err[0]
 
-
-@pytest.mark.skipif(not clipboard.can_clip, reason="Pyperclip could not find a copy/paste mechanism for your system")
+try:
+    # try getting the contents of the clipboard
+    _ = clipboard.get_paste_buffer()
+    # pyperclip raises at least the following types of exceptions
+    #   FileNotFoundError on Windows Subsystem for Linux (WSL) when Windows paths are removed from $PATH
+    #   ValueError for headless Linux systems without Gtk installed
+    #   AssertionError can be raised by paste_klipper().
+    #   PyperclipException for pyperclip-specific exceptions
+except Exception:
+    can_paste = False
+else:
+    can_paste = True
+
+@pytest.mark.skipif(can_paste, reason="Pyperclip could not find a copy/paste mechanism for your system")
 def test_send_to_paste_buffer(base_app):
     # Test writing to the PasteBuffer/Clipboard
     run_cmd(base_app, 'help >')
@@ -743,6 +755,35 @@ def test_send_to_paste_buffer(base_app):
     assert appended_contents.startswith(paste_contents)
     assert len(appended_contents) > len(paste_contents)
 
+def test_get_paste_buffer_exception(base_app, mocker, capsys):
+    # Force get_paste_buffer to throw an exception
+    pastemock = mocker.patch('pyperclip.paste')
+    pastemock.side_effect = ValueError('foo')
+
+    # Redirect command output to the clipboard
+    base_app.onecmd_plus_hooks('help > ')
+
+    # Make sure we got the exception output
+    out, err = capsys.readouterr()
+    assert out == ''
+    # this just checks that cmd2 is surfacing whatever error gets raised by pyperclip.paste
+    assert 'ValueError' in err and 'foo' in err
+
+def test_allow_clipboard_initializer(base_app):
+    assert base_app.allow_clipboard == True
+    noclipcmd = cmd2.Cmd(allow_clipboard=False)
+    assert noclipcmd.allow_clipboard == False
+
+# if clipboard access is not allowed, cmd2 should check that first
+# before it tries to do anything with pyperclip, that's why we can
+# safely run this test without skipping it if pyperclip doesn't
+# work in the test environment, like we do for test_send_to_paste_buffer()
+def test_allow_clipboard(base_app):
+    base_app.allow_clipboard = False
+    out, err = run_cmd(base_app, 'help >')
+    assert not out
+    assert "Clipboard access not allowed" in err
+
 
 def test_base_timing(base_app):
     base_app.feedback_to_output = False
@@ -1566,19 +1607,6 @@ def test_multiline_input_line_to_statement(multiline_app):
     assert statement.multiline_command == 'orate'
 
 
-def test_clipboard_failure(base_app, capsys):
-    # Force cmd2 clipboard to be disabled
-    base_app._can_clip = False
-
-    # Redirect command output to the clipboard when a clipboard isn't present
-    base_app.onecmd_plus_hooks('help > ')
-
-    # Make sure we got the error output
-    out, err = capsys.readouterr()
-    assert out == ''
-    assert 'Cannot redirect to paste buffer;' in err and 'pyperclip' in err
-
-
 class CommandResultApp(cmd2.Cmd):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)