command.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. # SPDX-License-Identifier: GPL-2.0+
  2. # Copyright (c) 2011 The Chromium OS Authors.
  3. #
  4. import os
  5. 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. # This permits interception of RunPipe for test purposes. If it is set to
  29. # a function, then that function is called with the pipe list being
  30. # executed. Otherwise, it is assumed to be a CommandResult object, and is
  31. # returned as the result for every RunPipe() call.
  32. # When this value is None, commands are executed as normal.
  33. test_result = None
  34. def RunPipe(pipe_list, infile=None, outfile=None,
  35. capture=False, capture_stderr=False, oneline=False,
  36. raise_on_error=True, cwd=None, **kwargs):
  37. """
  38. Perform a command pipeline, with optional input/output filenames.
  39. Args:
  40. pipe_list: List of command lines to execute. Each command line is
  41. piped into the next, and is itself a list of strings. For
  42. example [ ['ls', '.git'] ['wc'] ] will pipe the output of
  43. 'ls .git' into 'wc'.
  44. infile: File to provide stdin to the pipeline
  45. outfile: File to store stdout
  46. capture: True to capture output
  47. capture_stderr: True to capture stderr
  48. oneline: True to strip newline chars from output
  49. kwargs: Additional keyword arguments to cros_subprocess.Popen()
  50. Returns:
  51. CommandResult object
  52. """
  53. if test_result:
  54. if hasattr(test_result, '__call__'):
  55. result = test_result(pipe_list=pipe_list)
  56. if result:
  57. return result
  58. else:
  59. return test_result
  60. # No result: fall through to normal processing
  61. result = CommandResult()
  62. last_pipe = None
  63. pipeline = list(pipe_list)
  64. user_pipestr = '|'.join([' '.join(pipe) for pipe in pipe_list])
  65. kwargs['stdout'] = None
  66. kwargs['stderr'] = None
  67. while pipeline:
  68. cmd = pipeline.pop(0)
  69. if last_pipe is not None:
  70. kwargs['stdin'] = last_pipe.stdout
  71. elif infile:
  72. kwargs['stdin'] = open(infile, 'rb')
  73. if pipeline or capture:
  74. kwargs['stdout'] = cros_subprocess.PIPE
  75. elif outfile:
  76. kwargs['stdout'] = open(outfile, 'wb')
  77. if capture_stderr:
  78. kwargs['stderr'] = cros_subprocess.PIPE
  79. try:
  80. last_pipe = cros_subprocess.Popen(cmd, cwd=cwd, **kwargs)
  81. except Exception as err:
  82. result.exception = err
  83. if raise_on_error:
  84. raise Exception("Error running '%s': %s" % (user_pipestr, str))
  85. result.return_code = 255
  86. return result
  87. if capture:
  88. result.stdout, result.stderr, result.combined = (
  89. last_pipe.CommunicateFilter(None))
  90. if result.stdout and oneline:
  91. result.output = result.stdout.rstrip('\r\n')
  92. result.return_code = last_pipe.wait()
  93. else:
  94. result.return_code = os.waitpid(last_pipe.pid, 0)[1]
  95. if raise_on_error and result.return_code:
  96. raise Exception("Error running '%s'" % user_pipestr)
  97. return result
  98. def Output(*cmd, **kwargs):
  99. raise_on_error = kwargs.get('raise_on_error', True)
  100. return RunPipe([cmd], capture=True, raise_on_error=raise_on_error).stdout
  101. def OutputOneLine(*cmd, **kwargs):
  102. raise_on_error = kwargs.pop('raise_on_error', True)
  103. return (RunPipe([cmd], capture=True, oneline=True,
  104. raise_on_error=raise_on_error,
  105. **kwargs).stdout.strip())
  106. def Run(*cmd, **kwargs):
  107. return RunPipe([cmd], **kwargs).stdout
  108. def RunList(cmd):
  109. return RunPipe([cmd], capture=True).stdout
  110. def StopAll():
  111. cros_subprocess.stay_alive = False