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