| # Copyright (c) 2012 The Chromium OS Authors. | |
| # Use of this source code is governed by a BSD-style license that can be | |
| # found in the LICENSE file. | |
| # | |
| # Copyright (c) 2003-2005 by Peter Astrand <astrand@lysator.liu.se> | |
| # Licensed to PSF under a Contributor Agreement. | |
| # See http://www.python.org/2.4/license for licensing details. | |
| """Subprocress execution | |
| This module holds a subclass of subprocess.Popen with our own required | |
| features, mainly that we get access to the subprocess output while it | |
| is running rather than just at the end. This makes it easiler to show | |
| progress information and filter output in real time. | |
| """ | |
| import errno | |
| import os | |
| import pty | |
| import select | |
| import subprocess | |
| import sys | |
| import unittest | |
| # Import these here so the caller does not need to import subprocess also. | |
| PIPE = subprocess.PIPE | |
| STDOUT = subprocess.STDOUT | |
| PIPE_PTY = -3 # Pipe output through a pty | |
| stay_alive = True | |
| class Popen(subprocess.Popen): | |
| """Like subprocess.Popen with ptys and incremental output | |
| This class deals with running a child process and filtering its output on | |
| both stdout and stderr while it is running. We do this so we can monitor | |
| progress, and possibly relay the output to the user if requested. | |
| The class is similar to subprocess.Popen, the equivalent is something like: | |
| Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) | |
| But this class has many fewer features, and two enhancement: | |
| 1. Rather than getting the output data only at the end, this class sends it | |
| to a provided operation as it arrives. | |
| 2. We use pseudo terminals so that the child will hopefully flush its output | |
| to us as soon as it is produced, rather than waiting for the end of a | |
| line. | |
| Use CommunicateFilter() to handle output from the subprocess. | |
| """ | |
| def __init__(self, args, stdin=None, stdout=PIPE_PTY, stderr=PIPE_PTY, | |
| shell=False, cwd=None, env=None, **kwargs): | |
| """Cut-down constructor | |
| Args: | |
| args: Program and arguments for subprocess to execute. | |
| stdin: See subprocess.Popen() | |
| stdout: See subprocess.Popen(), except that we support the sentinel | |
| value of cros_subprocess.PIPE_PTY. | |
| stderr: See subprocess.Popen(), except that we support the sentinel | |
| value of cros_subprocess.PIPE_PTY. | |
| shell: See subprocess.Popen() | |
| cwd: Working directory to change to for subprocess, or None if none. | |
| env: Environment to use for this subprocess, or None to inherit parent. | |
| kwargs: No other arguments are supported at the moment. Passing other | |
| arguments will cause a ValueError to be raised. | |
| """ | |
| stdout_pty = None | |
| stderr_pty = None | |
| if stdout == PIPE_PTY: | |
| stdout_pty = pty.openpty() | |
| stdout = os.fdopen(stdout_pty[1]) | |
| if stderr == PIPE_PTY: | |
| stderr_pty = pty.openpty() | |
| stderr = os.fdopen(stderr_pty[1]) | |
| super(Popen, self).__init__(args, stdin=stdin, | |
| stdout=stdout, stderr=stderr, shell=shell, cwd=cwd, env=env, | |
| **kwargs) | |
| # If we're on a PTY, we passed the slave half of the PTY to the subprocess. | |
| # We want to use the master half on our end from now on. Setting this here | |
| # does make some assumptions about the implementation of subprocess, but | |
| # those assumptions are pretty minor. | |
| # Note that if stderr is STDOUT, then self.stderr will be set to None by | |
| # this constructor. | |
| if stdout_pty is not None: | |
| self.stdout = os.fdopen(stdout_pty[0]) | |
| if stderr_pty is not None: | |
| self.stderr = os.fdopen(stderr_pty[0]) | |
| # Insist that unit tests exist for other arguments we don't support. | |
| if kwargs: | |
| raise ValueError("Unit tests do not test extra args - please add tests") | |
| def CommunicateFilter(self, output): | |
| """Interact with process: Read data from stdout and stderr. | |
| This method runs until end-of-file is reached, then waits for the | |
| subprocess to terminate. | |
| The output function is sent all output from the subprocess and must be | |
| defined like this: | |
| def Output([self,] stream, data) | |
| Args: | |
| stream: the stream the output was received on, which will be | |
| sys.stdout or sys.stderr. | |
| data: a string containing the data | |
| Note: The data read is buffered in memory, so do not use this | |
| method if the data size is large or unlimited. | |
| Args: | |
| output: Function to call with each fragment of output. | |
| Returns: | |
| A tuple (stdout, stderr, combined) which is the data received on | |
| stdout, stderr and the combined data (interleaved stdout and stderr). | |
| Note that the interleaved output will only be sensible if you have | |
| set both stdout and stderr to PIPE or PIPE_PTY. Even then it depends on | |
| the timing of the output in the subprocess. If a subprocess flips | |
| between stdout and stderr quickly in succession, by the time we come to | |
| read the output from each we may see several lines in each, and will read | |
| all the stdout lines, then all the stderr lines. So the interleaving | |
| may not be correct. In this case you might want to pass | |
| stderr=cros_subprocess.STDOUT to the constructor. | |
| This feature is still useful for subprocesses where stderr is | |
| rarely used and indicates an error. | |
| Note also that if you set stderr to STDOUT, then stderr will be empty | |
| and the combined output will just be the same as stdout. | |
| """ | |
| read_set = [] | |
| write_set = [] | |
| stdout = None # Return | |
| stderr = None # Return | |
| if self.stdin: | |
| # Flush stdio buffer. This might block, if the user has | |
| # been writing to .stdin in an uncontrolled fashion. | |
| self.stdin.flush() | |
| if input: | |
| write_set.append(self.stdin) | |
| else: | |
| self.stdin.close() | |
| if self.stdout: | |
| read_set.append(self.stdout) | |
| stdout = [] | |
| if self.stderr and self.stderr != self.stdout: | |
| read_set.append(self.stderr) | |
| stderr = [] | |
| combined = [] | |
| input_offset = 0 | |
| while read_set or write_set: | |
| try: | |
| rlist, wlist, _ = select.select(read_set, write_set, [], 0.2) | |
| except select.error, e: | |
| if e.args[0] == errno.EINTR: | |
| continue | |
| raise | |
| if not stay_alive: | |
| self.terminate() | |
| if self.stdin in wlist: | |
| # When select has indicated that the file is writable, | |
| # we can write up to PIPE_BUF bytes without risk | |
| # blocking. POSIX defines PIPE_BUF >= 512 | |
| chunk = input[input_offset : input_offset + 512] | |
| bytes_written = os.write(self.stdin.fileno(), chunk) | |
| input_offset += bytes_written | |
| if input_offset >= len(input): | |
| self.stdin.close() | |
| write_set.remove(self.stdin) | |
| if self.stdout in rlist: | |
| data = "" | |
| # We will get an error on read if the pty is closed | |
| try: | |
| data = os.read(self.stdout.fileno(), 1024) | |
| except OSError: | |
| pass | |
| if data == "": | |
| self.stdout.close() | |
| read_set.remove(self.stdout) | |
| else: | |
| stdout.append(data) | |
| combined.append(data) | |
| if output: | |
| output(sys.stdout, data) | |
| if self.stderr in rlist: | |
| data = "" | |
| # We will get an error on read if the pty is closed | |
| try: | |
| data = os.read(self.stderr.fileno(), 1024) | |
| except OSError: | |
| pass | |
| if data == "": | |
| self.stderr.close() | |
| read_set.remove(self.stderr) | |
| else: | |
| stderr.append(data) | |
| combined.append(data) | |
| if output: | |
| output(sys.stderr, data) | |
| # All data exchanged. Translate lists into strings. | |
| if stdout is not None: | |
| stdout = ''.join(stdout) | |
| else: | |
| stdout = '' | |
| if stderr is not None: | |
| stderr = ''.join(stderr) | |
| else: | |
| stderr = '' | |
| combined = ''.join(combined) | |
| # Translate newlines, if requested. We cannot let the file | |
| # object do the translation: It is based on stdio, which is | |
| # impossible to combine with select (unless forcing no | |
| # buffering). | |
| if self.universal_newlines and hasattr(file, 'newlines'): | |
| if stdout: | |
| stdout = self._translate_newlines(stdout) | |
| if stderr: | |
| stderr = self._translate_newlines(stderr) | |
| self.wait() | |
| return (stdout, stderr, combined) | |
| # Just being a unittest.TestCase gives us 14 public methods. Unless we | |
| # disable this, we can only have 6 tests in a TestCase. That's not enough. | |
| # | |
| # pylint: disable=R0904 | |
| class TestSubprocess(unittest.TestCase): | |
| """Our simple unit test for this module""" | |
| class MyOperation: | |
| """Provides a operation that we can pass to Popen""" | |
| def __init__(self, input_to_send=None): | |
| """Constructor to set up the operation and possible input. | |
| Args: | |
| input_to_send: a text string to send when we first get input. We will | |
| add \r\n to the string. | |
| """ | |
| self.stdout_data = '' | |
| self.stderr_data = '' | |
| self.combined_data = '' | |
| self.stdin_pipe = None | |
| self._input_to_send = input_to_send | |
| if input_to_send: | |
| pipe = os.pipe() | |
| self.stdin_read_pipe = pipe[0] | |
| self._stdin_write_pipe = os.fdopen(pipe[1], 'w') | |
| def Output(self, stream, data): | |
| """Output handler for Popen. Stores the data for later comparison""" | |
| if stream == sys.stdout: | |
| self.stdout_data += data | |
| if stream == sys.stderr: | |
| self.stderr_data += data | |
| self.combined_data += data | |
| # Output the input string if we have one. | |
| if self._input_to_send: | |
| self._stdin_write_pipe.write(self._input_to_send + '\r\n') | |
| self._stdin_write_pipe.flush() | |
| def _BasicCheck(self, plist, oper): | |
| """Basic checks that the output looks sane.""" | |
| self.assertEqual(plist[0], oper.stdout_data) | |
| self.assertEqual(plist[1], oper.stderr_data) | |
| self.assertEqual(plist[2], oper.combined_data) | |
| # The total length of stdout and stderr should equal the combined length | |
| self.assertEqual(len(plist[0]) + len(plist[1]), len(plist[2])) | |
| def test_simple(self): | |
| """Simple redirection: Get process list""" | |
| oper = TestSubprocess.MyOperation() | |
| plist = Popen(['ps']).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| def test_stderr(self): | |
| """Check stdout and stderr""" | |
| oper = TestSubprocess.MyOperation() | |
| cmd = 'echo fred >/dev/stderr && false || echo bad' | |
| plist = Popen([cmd], shell=True).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(plist [0], 'bad\r\n') | |
| self.assertEqual(plist [1], 'fred\r\n') | |
| def test_shell(self): | |
| """Check with and without shell works""" | |
| oper = TestSubprocess.MyOperation() | |
| cmd = 'echo test >/dev/stderr' | |
| self.assertRaises(OSError, Popen, [cmd], shell=False) | |
| plist = Popen([cmd], shell=True).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(len(plist [0]), 0) | |
| self.assertEqual(plist [1], 'test\r\n') | |
| def test_list_args(self): | |
| """Check with and without shell works using list arguments""" | |
| oper = TestSubprocess.MyOperation() | |
| cmd = ['echo', 'test', '>/dev/stderr'] | |
| plist = Popen(cmd, shell=False).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(plist [0], ' '.join(cmd[1:]) + '\r\n') | |
| self.assertEqual(len(plist [1]), 0) | |
| oper = TestSubprocess.MyOperation() | |
| # this should be interpreted as 'echo' with the other args dropped | |
| cmd = ['echo', 'test', '>/dev/stderr'] | |
| plist = Popen(cmd, shell=True).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(plist [0], '\r\n') | |
| def test_cwd(self): | |
| """Check we can change directory""" | |
| for shell in (False, True): | |
| oper = TestSubprocess.MyOperation() | |
| plist = Popen('pwd', shell=shell, cwd='/tmp').CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(plist [0], '/tmp\r\n') | |
| def test_env(self): | |
| """Check we can change environment""" | |
| for add in (False, True): | |
| oper = TestSubprocess.MyOperation() | |
| env = os.environ | |
| if add: | |
| env ['FRED'] = 'fred' | |
| cmd = 'echo $FRED' | |
| plist = Popen(cmd, shell=True, env=env).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(plist [0], add and 'fred\r\n' or '\r\n') | |
| def test_extra_args(self): | |
| """Check we can't add extra arguments""" | |
| self.assertRaises(ValueError, Popen, 'true', close_fds=False) | |
| def test_basic_input(self): | |
| """Check that incremental input works | |
| We set up a subprocess which will prompt for name. When we see this prompt | |
| we send the name as input to the process. It should then print the name | |
| properly to stdout. | |
| """ | |
| oper = TestSubprocess.MyOperation('Flash') | |
| prompt = 'What is your name?: ' | |
| cmd = 'echo -n "%s"; read name; echo Hello $name' % prompt | |
| plist = Popen([cmd], stdin=oper.stdin_read_pipe, | |
| shell=True).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(len(plist [1]), 0) | |
| self.assertEqual(plist [0], prompt + 'Hello Flash\r\r\n') | |
| def test_isatty(self): | |
| """Check that ptys appear as terminals to the subprocess""" | |
| oper = TestSubprocess.MyOperation() | |
| cmd = ('if [ -t %d ]; then echo "terminal %d" >&%d; ' | |
| 'else echo "not %d" >&%d; fi;') | |
| both_cmds = '' | |
| for fd in (1, 2): | |
| both_cmds += cmd % (fd, fd, fd, fd, fd) | |
| plist = Popen(both_cmds, shell=True).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(plist [0], 'terminal 1\r\n') | |
| self.assertEqual(plist [1], 'terminal 2\r\n') | |
| # Now try with PIPE and make sure it is not a terminal | |
| oper = TestSubprocess.MyOperation() | |
| plist = Popen(both_cmds, stdout=subprocess.PIPE, stderr=subprocess.PIPE, | |
| shell=True).CommunicateFilter(oper.Output) | |
| self._BasicCheck(plist, oper) | |
| self.assertEqual(plist [0], 'not 1\n') | |
| self.assertEqual(plist [1], 'not 2\n') | |
| if __name__ == '__main__': | |
| unittest.main() |