command.py 4.2 KB

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