stgit/run.py

   1 import datetime
   2 import io
   3 import os
   4 import subprocess
   5
   6 from stgit.compat import environ_get, fsencode_utf8
   7 from stgit.exception import StgException
   8 from stgit.out import MessagePrinter, out
   9
  10 __copyright__ = """
  11 Copyright (C) 2007, Karl Hasselström <kha@treskal.com>
  12
  13 This program is free software; you can redistribute it and/or modify
  14 it under the terms of the GNU General Public License version 2 as
  15 published by the Free Software Foundation.
  16
  17 This program is distributed in the hope that it will be useful,
  18 but WITHOUT ANY WARRANTY; without even the implied warranty of
  19 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  20 GNU General Public License for more details.
  21
  22 You should have received a copy of the GNU General Public License
  23 along with this program; if not, see http://www.gnu.org/licenses/.
  24 """
  25
  26
  27 class RunException(StgException):
  28     """Exception raised for subprocess failures."""
  29
  30
  31 def get_log_mode(spec):
  32     if ':' not in spec:
  33         spec += ':'
  34     (log_mode, outfile) = spec.split(':', 1)
  35     all_log_modes = ['debug', 'profile']
  36     if log_mode and log_mode not in all_log_modes:
  37         out.warn(
  38             ('Unknown log mode "%s" specified in $STGIT_SUBPROCESS_LOG.' % log_mode),
  39             'Valid values are: %s' % ', '.join(all_log_modes),
  40         )
  41     if outfile:
  42         f = MessagePrinter(io.open(outfile, 'a', encoding='utf-8'))
  43     else:
  44         f = out
  45     return (log_mode, f)
  46
  47
  48 _log_mode, _logfile = get_log_mode(environ_get('STGIT_SUBPROCESS_LOG', ''))
  49 if _log_mode == 'profile':
  50     _log_starttime = datetime.datetime.now()
  51     _log_subproctime = 0.0
  52
  53
  54 def duration(t1, t2):
  55     d = t2 - t1
  56     return 86400 * d.days + d.seconds + 1e-6 * d.microseconds
  57
  58
  59 def finish_logging():
  60     if _log_mode != 'profile':
  61         return
  62     ttime = duration(_log_starttime, datetime.datetime.now())
  63     rtime = ttime - _log_subproctime
  64     _logfile.info(
  65         'Total time: %1.3f s' % ttime,
  66         'Time spent in subprocess calls: %1.3f s (%1.1f%%)'
  67         % (_log_subproctime, 100 * _log_subproctime / ttime),
  68         'Remaining time: %1.3f s (%1.1f%%)' % (rtime, 100 * rtime / ttime),
  69     )
  70
  71
  72 class Run:
  73     exc = RunException
  74
  75     def __init__(self, *cmd):
  76         self._cmd = list(cmd)
  77         self._good_retvals = [0]
  78         self._env = self._cwd = None
  79         self._indata = None
  80         self._in_encoding = 'utf-8'
  81         self._out_encoding = 'utf-8'
  82         self._discard_stderr = False
  83
  84     def _prep_cmd(self):
  85         return [fsencode_utf8(c) for c in self._cmd]
  86
  87     def _prep_env(self):
  88         # Windows requires a dict of strings as env parameter, so don't encode for Windows
  89         if self._env and os.name != 'nt':
  90             return {fsencode_utf8(k): fsencode_utf8(v) for k, v in self._env.items()}
  91         else:
  92             return self._env
  93
  94     def _log_start(self):
  95         if _log_mode == 'debug':
  96             _logfile.start('Running subprocess %s' % self._cmd)
  97             if self._cwd is not None:
  98                 _logfile.info('cwd: %s' % self._cwd)
  99             if self._env is not None:
 100                 for k in sorted(self._env):
 101                     v = environ_get(k)
 102                     if v is None or v != self._env[k]:
 103                         _logfile.info('%s: %s' % (k, self._env[k]))
 104         elif _log_mode == 'profile':
 105             _logfile.start('Running subprocess %s' % self._cmd)
 106             self._starttime = datetime.datetime.now()
 107
 108     def _log_end(self, retcode):
 109         global _log_subproctime, _log_starttime
 110         if _log_mode == 'debug':
 111             _logfile.done('return code: %d' % retcode)
 112         elif _log_mode == 'profile':
 113             n = datetime.datetime.now()
 114             d = duration(self._starttime, n)
 115             _logfile.done('%1.3f s' % d)
 116             _log_subproctime += d
 117             _logfile.info(
 118                 'Time since program start: %1.3f s' % duration(_log_starttime, n)
 119             )
 120
 121     def _check_exitcode(self):
 122         if self._good_retvals is None:
 123             return
 124         if self.exitcode not in self._good_retvals:
 125             raise self.exc('%s failed with code %d' % (self._cmd[0], self.exitcode))
 126
 127     def _run_io(self):
 128         """Run with captured IO."""
 129         self._log_start()
 130         try:
 131             p = subprocess.Popen(
 132                 self._prep_cmd(),
 133                 env=self._prep_env(),
 134                 cwd=self._cwd,
 135                 stdin=subprocess.PIPE,
 136                 stdout=subprocess.PIPE,
 137                 stderr=subprocess.PIPE,
 138             )
 139             outdata, errdata = p.communicate(self._indata)
 140             self.exitcode = p.returncode
 141         except OSError as e:
 142             raise self.exc('%s failed: %s' % (self._cmd[0], e))
 143         if errdata and not self._discard_stderr:
 144             out.err_bytes(errdata)
 145         self._log_end(self.exitcode)
 146         self._check_exitcode()
 147         if self._out_encoding:
 148             return outdata.decode(self._out_encoding)
 149         else:
 150             return outdata
 151
 152     def _run_noio(self):
 153         """Run without captured IO."""
 154         assert self._indata is None
 155         self._log_start()
 156         try:
 157             p = subprocess.Popen(
 158                 self._prep_cmd(),
 159                 env=self._prep_env(),
 160                 cwd=self._cwd,
 161             )
 162             self.exitcode = p.wait()
 163         except OSError as e:
 164             raise self.exc('%s failed: %s' % (self._cmd[0], e))
 165         self._log_end(self.exitcode)
 166         self._check_exitcode()
 167
 168     def run_background(self):
 169         """Run as a background process."""
 170         assert self._indata is None
 171         assert self._in_encoding is None
 172         assert self._out_encoding is None
 173         try:
 174             return subprocess.Popen(
 175                 self._prep_cmd(),
 176                 env=self._prep_env(),
 177                 cwd=self._cwd,
 178                 stdin=subprocess.PIPE,
 179                 stdout=subprocess.PIPE,
 180                 stderr=subprocess.PIPE,
 181             )
 182         except OSError as e:
 183             raise self.exc('%s failed: %s' % (self._cmd[0], e))
 184
 185     def returns(self, retvals):
 186         self._good_retvals = retvals
 187         return self
 188
 189     def discard_exitcode(self):
 190         self._good_retvals = None
 191         return self
 192
 193     def discard_stderr(self, discard=True):
 194         self._discard_stderr = discard
 195         return self
 196
 197     def env(self, env):
 198         self._env = os.environ.copy()
 199         self._env.update(env)
 200         return self
 201
 202     def cwd(self, cwd):
 203         self._cwd = cwd
 204         return self
 205
 206     def encoding(self, encoding):
 207         self._in_encoding = encoding
 208         return self
 209
 210     def decoding(self, encoding):
 211         self._out_encoding = encoding
 212         return self
 213
 214     def raw_input(self, indata):
 215         if self._in_encoding:
 216             self._indata = indata.encode(self._in_encoding)
 217         else:
 218             self._indata = indata
 219         return self
 220
 221     def input_nulterm(self, lines):
 222         return self.raw_input('\0'.join(lines))
 223
 224     def no_output(self):
 225         outdata = self._run_io()
 226         if outdata:
 227             raise self.exc('%s produced output' % self._cmd[0])
 228
 229     def discard_output(self):
 230         self._run_io()
 231
 232     def raw_output(self):
 233         return self._run_io()
 234
 235     def output_lines(self, sep='\n'):
 236         outdata = self._run_io()
 237         if outdata.endswith(sep):
 238             outdata = outdata[:-1]
 239         if outdata:
 240             return outdata.split(sep)
 241         else:
 242             return []
 243
 244     def output_one_line(self, sep='\n'):
 245         outlines = self.output_lines(sep)
 246         if len(outlines) == 1:
 247             return outlines[0]
 248         else:
 249             raise self.exc(
 250                 '%s produced %d lines, expected 1' % (self._cmd[0], len(outlines))
 251             )
 252
 253     def run(self):
 254         """Just run, with no IO redirection."""
 255         self._run_noio()