command.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. # SPDX-License-Identifier: GPL-2.0+
  2. # Copyright (c) 2011 The Chromium OS Authors.
  3. #
  4. import os
  5. from u_boot_pylib 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, stdout='', stderr='', combined='', return_code=0,
  16. exception=None):
  17. self.stdout = stdout
  18. self.stderr = stderr
  19. self.combined = combined
  20. self.return_code = return_code
  21. self.exception = exception
  22. def to_output(self, binary):
  23. if not binary:
  24. self.stdout = self.stdout.decode('utf-8')
  25. self.stderr = self.stderr.decode('utf-8')
  26. self.combined = self.combined.decode('utf-8')
  27. return self
  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 run_pipe() call.
  32. # When this value is None, commands are executed as normal.
  33. test_result = None
  34. def run_pipe(pipe_list, infile=None, outfile=None,
  35. capture=False, capture_stderr=False, oneline=False,
  36. raise_on_error=True, cwd=None, binary=False,
  37. output_func=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. output_func: Output function to call with each output fragment
  51. (if it returns True the function terminates)
  52. kwargs: Additional keyword arguments to cros_subprocess.Popen()
  53. Returns:
  54. CommandResult object
  55. """
  56. if test_result:
  57. if hasattr(test_result, '__call__'):
  58. # pylint: disable=E1102
  59. result = test_result(pipe_list=pipe_list)
  60. if result:
  61. return result
  62. else:
  63. return test_result
  64. # No result: fall through to normal processing
  65. result = CommandResult(b'', b'', b'')
  66. last_pipe = None
  67. pipeline = list(pipe_list)
  68. user_pipestr = '|'.join([' '.join(pipe) for pipe in pipe_list])
  69. kwargs['stdout'] = None
  70. kwargs['stderr'] = None
  71. while pipeline:
  72. cmd = pipeline.pop(0)
  73. if last_pipe is not None:
  74. kwargs['stdin'] = last_pipe.stdout
  75. elif infile:
  76. kwargs['stdin'] = open(infile, 'rb')
  77. if pipeline or capture:
  78. kwargs['stdout'] = cros_subprocess.PIPE
  79. elif outfile:
  80. kwargs['stdout'] = open(outfile, 'wb')
  81. if capture_stderr:
  82. kwargs['stderr'] = cros_subprocess.PIPE
  83. try:
  84. last_pipe = cros_subprocess.Popen(cmd, cwd=cwd, **kwargs)
  85. except Exception as err:
  86. result.exception = err
  87. if raise_on_error:
  88. raise Exception("Error running '%s': %s" % (user_pipestr, str))
  89. result.return_code = 255
  90. return result.to_output(binary)
  91. if capture:
  92. result.stdout, result.stderr, result.combined = (
  93. last_pipe.communicate_filter(output_func))
  94. if result.stdout and oneline:
  95. result.output = result.stdout.rstrip(b'\r\n')
  96. result.return_code = last_pipe.wait()
  97. else:
  98. result.return_code = os.waitpid(last_pipe.pid, 0)[1]
  99. if raise_on_error and result.return_code:
  100. raise Exception("Error running '%s'" % user_pipestr)
  101. return result.to_output(binary)
  102. def output(*cmd, **kwargs):
  103. kwargs['raise_on_error'] = kwargs.get('raise_on_error', True)
  104. return run_pipe([cmd], capture=True, **kwargs).stdout
  105. def output_one_line(*cmd, **kwargs):
  106. """Run a command and output it as a single-line string
  107. The command us expected to produce a single line of output
  108. Returns:
  109. String containing output of command
  110. """
  111. raise_on_error = kwargs.pop('raise_on_error', True)
  112. result = run_pipe([cmd], capture=True, oneline=True,
  113. raise_on_error=raise_on_error, **kwargs).stdout.strip()
  114. return result
  115. def run(*cmd, **kwargs):
  116. return run_pipe([cmd], **kwargs).stdout
  117. def run_list(cmd):
  118. return run_pipe([cmd], capture=True).stdout
  119. def stop_all():
  120. cros_subprocess.stay_alive = False