third_party/pigweed/repo/pw_cli/py/pw_cli/process.py

   1 # Copyright 2019 The Pigweed Authors
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
   4 # use this file except in compliance with the License. You may obtain a copy of
   5 # the License at
   6 #
   7 #     https://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12 # License for the specific language governing permissions and limitations under
  13 # the License.
  14 """Module for running subprocesses from pw and capturing their output."""
  15
  16 import asyncio
  17 import logging
  18 import os
  19 import shlex
  20 import tempfile
  21 from typing import IO, Tuple, Union
  22
  23 import pw_cli.color
  24 import pw_cli.log
  25
  26 _COLOR = pw_cli.color.colors()
  27 _LOG = logging.getLogger(__name__)
  28
  29 # Environment variable passed down to subprocesses to indicate that they are
  30 # running as a subprocess. Can be imported by other code.
  31 PW_SUBPROCESS_ENV = 'PW_SUBPROCESS'
  32
  33
  34 class CompletedProcess:
  35     """Information about a process executed in run_async."""
  36     def __init__(self, process: 'asyncio.subprocess.Process',
  37                  output: Union[bytes, IO[bytes]]):
  38         assert process.returncode is not None
  39         self.returncode: int = process.returncode
  40         self.pid = process.pid
  41         self._output = output
  42
  43     @property
  44     def output(self) -> bytes:
  45         # If the output is a file, read it, then close it.
  46         if not isinstance(self._output, bytes):
  47             with self._output as file:
  48                 file.flush()
  49                 file.seek(0)
  50                 self._output = file.read()
  51
  52         return self._output
  53
  54
  55 async def _run_and_log(program: str, args: Tuple[str, ...], env: dict):
  56     process = await asyncio.create_subprocess_exec(
  57         program,
  58         *args,
  59         stdout=asyncio.subprocess.PIPE,
  60         stderr=asyncio.subprocess.STDOUT,
  61         env=env)
  62
  63     output = bytearray()
  64
  65     if process.stdout:
  66         while True:
  67             line = await process.stdout.readline()
  68             if not line:
  69                 break
  70
  71             output += line
  72             _LOG.log(pw_cli.log.LOGLEVEL_STDOUT, '[%s] %s',
  73                      _COLOR.bold_white(process.pid),
  74                      line.decode(errors='replace').rstrip())
  75
  76     return process, bytes(output)
  77
  78
  79 async def run_async(program: str,
  80                     *args: str,
  81                     log_output: bool = False) -> CompletedProcess:
  82     """Runs a command, capturing and optionally logging its output.
  83
  84     Returns a CompletedProcess with details from the process.
  85     """
  86
  87     _LOG.debug('Running `%s`',
  88                ' '.join(shlex.quote(arg) for arg in [program, *args]))
  89
  90     env = os.environ.copy()
  91     env[PW_SUBPROCESS_ENV] = '1'
  92     output: Union[bytes, IO[bytes]]
  93
  94     if log_output:
  95         process, output = await _run_and_log(program, args, env)
  96     else:
  97         output = tempfile.TemporaryFile()
  98         process = await asyncio.create_subprocess_exec(
  99             program,
 100             *args,
 101             stdout=output,
 102             stderr=asyncio.subprocess.STDOUT,
 103             env=env)
 104
 105     if await process.wait() == 0:
 106         _LOG.info('%s exited successfully', program)
 107     else:
 108         _LOG.error('%s exited with status %d', program, process.returncode)
 109
 110     return CompletedProcess(process, output)
 111
 112
 113 def run(program: str, *args: str, **kwargs) -> CompletedProcess:
 114     """Synchronous wrapper for run_async."""
 115     return asyncio.run(run_async(program, *args, **kwargs))