Skip to content

Commit de9cda2

Browse files
committed
First attempt at improving how command output gets piped to a shell command
Now a real pipe is created to a subprocess. This has many advantages and should "just work" like intended with all commands. One downside is to work properly on Python 2.7, it requires the subprocess32 module which is the subprocess module from Python 3.2 backported to Python 2.7. Another downside, is that unit testing the feature is now more difficult. This still needs to be tested for compatibility across all OSes and supported versions of Python. The user needs to be careful if designing multi-threaded cmd2 applications that do command processing in other threads and those threads can make calls to self.stdout.write to put in a try/except to catch Broken Pipe errors which can occur for long running commands if the user closes the shell subprocess before the command is finished.
1 parent 11f14ed commit de9cda2

File tree

3 files changed

+89
-56
lines changed

3 files changed

+89
-56
lines changed

cmd2.py

Lines changed: 55 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,13 @@
3030
import collections
3131
import datetime
3232
import glob
33+
import io
3334
import optparse
3435
import os
3536
import platform
3637
import re
3738
import shlex
3839
import six
39-
import subprocess
4040
import sys
4141
import tempfile
4242
import traceback
@@ -59,6 +59,13 @@
5959
# itertools.zip() for Python 2 or zip() for Python 3 - produces an iterator in both cases
6060
from six.moves import zip
6161

62+
# If using Python 2.7, try to use the subprocess32 package backported from Python 3.2 due to various improvements
63+
# NOTE: The feature to pipe output to a shell command won't work correctly in Python 2.7 without this
64+
try:
65+
import subprocess32 as subprocess
66+
except ImportError:
67+
import subprocess
68+
6269
# Detect whether IPython is installed to determine if the built-in "ipy" command should be included
6370
ipython_available = True
6471
try:
@@ -513,9 +520,6 @@ def __init__(self, completekey='tab', stdin=None, stdout=None, use_ipython=False
513520
self.kept_state = None
514521
self.kept_sys = None
515522

516-
# Used for a temp file during a pipe (needed tempfile instead of real pipe for Python 3.x prior to 3.5)
517-
self._temp_filename = None
518-
519523
# Codes used for exit conditions
520524
self._STOP_AND_EXIT = True # cmd convention
521525

@@ -531,6 +535,9 @@ def __init__(self, completekey='tab', stdin=None, stdout=None, use_ipython=False
531535
# Used load command to store the current script dir as a LIFO queue to support _relative_load command
532536
self._script_dir = []
533537

538+
# Used when piping command output to a shell command
539+
self.pipe_proc = None
540+
534541
# ----- Methods related to presenting output to the user -----
535542

536543
@property
@@ -765,18 +772,29 @@ def _redirect_output(self, statement):
765772
"""
766773
if statement.parsed.pipeTo:
767774
self.kept_state = Statekeeper(self, ('stdout',))
768-
self.kept_sys = Statekeeper(sys, ('stdout',))
769-
sys.stdout = self.stdout
770775

771-
# NOTE: We couldn't get a real pipe working via subprocess for Python 3.x prior to 3.5.
772-
# So to allow compatibility with Python 2.7 and 3.3+ we are redirecting output to a temporary file.
773-
# And once command is complete we are the temp file as stdin for the shell command to pipe to.
774-
# TODO: Once support for Python 3.x prior to 3.5 is no longer necessary, replace with a real subprocess pipe
776+
# Create a pipe with read and write sides
777+
read_fd, write_fd = os.pipe()
775778

776-
# Redirect stdout to a temporary file
777-
fd, self._temp_filename = tempfile.mkstemp()
778-
os.close(fd)
779-
self.stdout = open(self._temp_filename, 'w')
779+
# Open each side of the pipe and set stdout accordingly
780+
# noinspection PyTypeChecker
781+
self.stdout = io.open(write_fd, 'w')
782+
# noinspection PyTypeChecker
783+
subproc_stdin = io.open(read_fd, 'r')
784+
785+
# If you don't set shell=True, subprocess failure will throw an exception
786+
try:
787+
self.pipe_proc = subprocess.Popen(shlex.split(statement.parsed.pipeTo), stdin=subproc_stdin)
788+
except Exception as ex:
789+
# Restore stdout to what it was and close the pipe
790+
self.stdout.close()
791+
subproc_stdin.close()
792+
self.pipe_proc = None
793+
self.kept_state.restore()
794+
self.kept_state = None
795+
796+
# Re-raise the exception
797+
raise ex
780798
elif statement.parsed.output:
781799
if (not statement.parsed.outputTo) and (not can_clip):
782800
raise EnvironmentError('Cannot redirect to paste buffer; install ``xclip`` and re-run to enable')
@@ -797,32 +815,29 @@ def _restore_output(self, statement):
797815
798816
:param statement: ParsedString - subclass of str which also contains pyparsing ParseResults instance
799817
"""
800-
if self.kept_state:
801-
try:
802-
if statement.parsed.output:
803-
if not statement.parsed.outputTo:
804-
self.stdout.seek(0)
805-
write_to_paste_buffer(self.stdout.read())
806-
finally:
807-
self.stdout.close()
808-
self.kept_state.restore()
809-
self.kept_sys.restore()
810-
self.kept_state = None
811-
812-
if statement.parsed.pipeTo:
813-
# Pipe the contents of tempfile to the specified shell command
814-
with open(self._temp_filename) as fd:
815-
pipe_proc = subprocess.Popen(shlex.split(statement.parsed.pipeTo), stdin=fd,
816-
stdout=subprocess.PIPE)
817-
output, _ = pipe_proc.communicate()
818-
819-
if six.PY3:
820-
self.stdout.write(output.decode())
821-
else:
822-
self.stdout.write(output)
823-
824-
os.remove(self._temp_filename)
825-
self._temp_filename = None
818+
# If we have redirected output to a file or the clipboard or piped it to a shell command, then restore state
819+
if self.kept_state is not None:
820+
# If we redirected output to the clipboard
821+
if statement.parsed.output and not statement.parsed.outputTo:
822+
self.stdout.seek(0)
823+
write_to_paste_buffer(self.stdout.read())
824+
825+
# Close the file or pipe that stdout was redirected to
826+
self.stdout.close()
827+
828+
# If we were piping output to a shell command, then close the subprocess the shell command was running in
829+
if self.pipe_proc is not None:
830+
self.pipe_proc.communicate()
831+
self.pipe_proc = None
832+
833+
# Restore self.stdout
834+
self.kept_state.restore()
835+
self.kept_state = None
836+
837+
# Restore sys.stdout if need be
838+
if self.kept_sys is not None:
839+
self.kept_sys.restore()
840+
self.kept_sys = None
826841

827842
def _func_named(self, arg):
828843
"""Gets the method name associated with a given command.

tests/test_cmd2.py

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -543,33 +543,50 @@ def test_input_redirection(base_app, request):
543543

544544
# NOTE: File 'redirect.txt" contains 1 word "history"
545545

546-
# Verify that redirecting input from a file works
546+
# Verify that redirecting input ffom a file works
547547
out = run_cmd(base_app, 'help < {}'.format(filename))
548548
expected = normalize(HELP_HISTORY)
549549
assert out == expected
550550

551551

552-
def test_pipe_to_shell(base_app):
552+
def test_pipe_to_shell(base_app, capsys):
553553
if sys.platform == "win32":
554554
# Windows
555+
command = 'help | sort'
555556
# Get help menu and pipe it's output to the sort shell command
556-
out = run_cmd(base_app, 'help | sort')
557-
expected = ['', '', '_relative_load edit history py quit save shell show',
558-
'========================================',
559-
'cmdenvironment help load pyscript run set shortcuts',
560-
'Documented commands (type help <topic>):']
561-
assert out == expected
557+
run_cmd(base_app, 'help | sort')
558+
# expected = ['', '', '_relative_load edit history py quit save shell show',
559+
# '========================================',
560+
# 'cmdenvironment help load pyscript run set shortcuts',
561+
# 'Documented commands (type help <topic>):']
562+
# assert out == expected
562563
else:
563564
# Mac and Linux
564565
# Get help on help and pipe it's output to the input of the word count shell command
565-
out = run_cmd(base_app, 'help help | wc')
566-
567-
# Mac and Linux wc behave the same when piped from shell, but differently when piped stdin from file directly
568-
if sys.platform == 'darwin':
569-
expected = normalize("1 11 70")
570-
else:
571-
expected = normalize("1 11 70")
572-
assert out[0].strip() == expected[0].strip()
566+
command = 'help help | wc'
567+
# # Mac and Linux wc behave the same when piped from shell, but differently when piped stdin from file directly
568+
# if sys.platform == 'darwin':
569+
# expected = "1 11 70"
570+
# else:
571+
# expected = "1 11 70"
572+
# assert out.strip() == expected.strip()
573+
574+
run_cmd(base_app, command)
575+
out, err = capsys.readouterr()
576+
577+
# Unfortunately with the improved way of piping output to a subprocess, there isn't any good way of getting
578+
# access to the output produced by that subprocess within a unit test, but we can verify that no error occured
579+
assert not err
580+
581+
582+
583+
def test_pipe_to_shell_error(base_app, capsys):
584+
# Try to pipe command output to a shell command that doesn't exist in order to produce an error
585+
run_cmd(base_app, 'help | foobarbaz.this_does_not_exist')
586+
out, err = capsys.readouterr()
587+
588+
assert not out
589+
assert err.startswith("EXCEPTION of type 'FileNotFoundError' occurred with message:")
573590

574591

575592
@pytest.mark.skipif(not cmd2.can_clip,

tox.ini

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ deps =
1919
pytest-cov
2020
pytest-xdist
2121
six
22+
subprocess32
2223
commands =
2324
py.test {posargs: -n 2} --cov=cmd2 --cov-report=term-missing
2425
codecov

0 commit comments

Comments
 (0)