command.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. # SPDX-License-Identifier: GPL-2.0+
  2. # Copyright (c) 2011 The Chromium OS Authors.
  3. #
  4. import os
  5. import cros_subprocess
  6. import tools
  7. """Shell command ease-ups for Python."""
  8. class CommandResult:
  9. """A class which captures the result of executing a command.
  10. Members:
  11. stdout: stdout obtained from command, as a string
  12. stderr: stderr obtained from command, as a string
  13. return_code: Return code from command
  14. exception: Exception received, or None if all ok
  15. """
  16. def __init__(self):
  17. self.stdout = None
  18. self.stderr = None
  19. self.combined = None
  20. self.return_code = None
  21. self.exception = None
  22. def __init__(self, stdout='', stderr='', combined='', return_code=0,
  23. exception=None):
  24. self.stdout = stdout
  25. self.stderr = stderr
  26. self.combined = combined
  27. self.return_code = return_code
  28. self.exception = exception
  29. def ToOutput(self, binary):
  30. if not binary:
  31. self.stdout = tools.ToString(self.stdout)
  32. self.stderr = tools.ToString(self.stderr)
  33. self.combined = tools.ToString(self.combined)
  34. return self
  35. # This permits interception of RunPipe for test purposes. If it is set to
  36. # a function, then that function is called with the pipe list being
  37. # executed. Otherwise, it is assumed to be a CommandResult object, and is
  38. # returned as the result for every RunPipe() call.
  39. # When this value is None, commands are executed as normal.
  40. test_result = None
  41. def RunPipe(pipe_list, infile=None, outfile=None,
  42. capture=False, capture_stderr=False, oneline=False,
  43. raise_on_error=True, cwd=None, binary=False, **kwargs):
  44. """
  45. Perform a command pipeline, with optional input/output filenames.
  46. Args:
  47. pipe_list: List of command lines to execute. Each command line is
  48. piped into the next, and is itself a list of strings. For
  49. example [ ['ls', '.git'] ['wc'] ] will pipe the output of
  50. 'ls .git' into 'wc'.
  51. infile: File to provide stdin to the pipeline
  52. outfile: File to store stdout
  53. capture: True to capture output
  54. capture_stderr: True to capture stderr
  55. oneline: True to strip newline chars from output
  56. kwargs: Additional keyword arguments to cros_subprocess.Popen()
  57. Returns:
  58. CommandResult object
  59. """
  60. if test_result:
  61. if hasattr(test_result, '__call__'):
  62. result = test_result(pipe_list=pipe_list)
  63. if result:
  64. return result
  65. else:
  66. return test_result
  67. # No result: fall through to normal processing
  68. result = CommandResult(b'', b'', b'')
  69. last_pipe = None
  70. pipeline = list(pipe_list)
  71. user_pipestr = '|'.join([' '.join(pipe) for pipe in pipe_list])
  72. kwargs['stdout'] = None
  73. kwargs['stderr'] = None
  74. while pipeline:
  75. cmd = pipeline.pop(0)
  76. if last_pipe is not None:
  77. kwargs['stdin'] = last_pipe.stdout
  78. elif infile:
  79. kwargs['stdin'] = open(infile, 'rb')
  80. if pipeline or capture:
  81. kwargs['stdout'] = cros_subprocess.PIPE
  82. elif outfile:
  83. kwargs['stdout'] = open(outfile, 'wb')
  84. if capture_stderr:
  85. kwargs['stderr'] = cros_subprocess.PIPE
  86. try:
  87. last_pipe = cros_subprocess.Popen(cmd, cwd=cwd, **kwargs)
  88. except Exception as err:
  89. result.exception = err
  90. if raise_on_error:
  91. raise Exception("Error running '%s': %s" % (user_pipestr, str))
  92. result.return_code = 255
  93. return result.ToOutput(binary)
  94. if capture:
  95. result.stdout, result.stderr, result.combined = (
  96. last_pipe.CommunicateFilter(None))
  97. if result.stdout and oneline:
  98. result.output = result.stdout.rstrip(b'\r\n')
  99. result.return_code = last_pipe.wait()
  100. else:
  101. result.return_code = os.waitpid(last_pipe.pid, 0)[1]
  102. if raise_on_error and result.return_code:
  103. raise Exception("Error running '%s'" % user_pipestr)
  104. return result.ToOutput(binary)
  105. def Output(*cmd, **kwargs):
  106. kwargs['raise_on_error'] = kwargs.get('raise_on_error', True)
  107. return RunPipe([cmd], capture=True, **kwargs).stdout
  108. def OutputOneLine(*cmd, **kwargs):
  109. """Run a command and output it as a single-line string
  110. The command us expected to produce a single line of output
  111. Returns:
  112. String containing output of command
  113. """
  114. raise_on_error = kwargs.pop('raise_on_error', True)
  115. result = RunPipe([cmd], capture=True, oneline=True,
  116. raise_on_error=raise_on_error, **kwargs).stdout.strip()
  117. return result
  118. def Run(*cmd, **kwargs):
  119. return RunPipe([cmd], **kwargs).stdout
  120. def RunList(cmd):
  121. return RunPipe([cmd], capture=True).stdout
  122. def StopAll():
  123. cros_subprocess.stay_alive = False