3030import collections
3131import datetime
3232import glob
33+ import io
3334import optparse
3435import os
3536import platform
3637import re
3738import shlex
3839import six
39- import subprocess
4040import sys
4141import tempfile
4242import traceback
5959# itertools.zip() for Python 2 or zip() for Python 3 - produces an iterator in both cases
6060from 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
6370ipython_available = True
6471try :
@@ -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.
0 commit comments