github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/acceptancetests/jujupy/timeout.py (about)

     1  #!/usr/bin/env python
     2  #
     3  # This file is part of JujuPy, a library for driving the Juju CLI.
     4  # Copyright 2015, 2017 Canonical Ltd.
     5  #
     6  # This program is free software: you can redistribute it and/or modify it
     7  # under the terms of the Lesser GNU General Public License version 3, as
     8  # published by the Free Software Foundation.
     9  #
    10  # This program is distributed in the hope that it will be useful, but WITHOUT
    11  # ANY WARRANTY; without even the implied warranties of MERCHANTABILITY,
    12  # SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR PURPOSE.  See the Lesser
    13  # GNU General Public License for more details.
    14  #
    15  # You should have received a copy of the Lesser GNU General Public License
    16  # along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17  
    18  
    19  """A Python implementation of the *nix utility for use on all platforms."""
    20  from argparse import ArgumentParser
    21  from itertools import chain
    22  import signal
    23  import subprocess
    24  import sys
    25  import time
    26  
    27  from utility import until_timeout
    28  
    29  
    30  # Generate a list of all signals that can be used with Popen.send_signal on
    31  # this platform.
    32  if sys.platform == 'win32':
    33      signals = {
    34          'TERM': signal.SIGTERM,
    35          # CTRL_C_EVENT is also supposed to work, but experience shows
    36          # otherwise.
    37          'CTRL_BREAK': signal.CTRL_BREAK_EVENT,
    38          }
    39  else:
    40      # Blech.  No equivalent of errno.errorcode for signals.
    41      signals = dict(
    42          (x[3:], getattr(signal, x)) for x in dir(signal) if
    43          x.startswith('SIG') and x not in ('SIG_DFL', 'SIG_IGN'))
    44  
    45  
    46  def parse_args(argv=None):
    47      parser = ArgumentParser()
    48      parser.add_argument('duration', type=float)
    49  
    50      parser.add_argument(
    51          '--signal', default='TERM', choices=sorted(signals.keys()))
    52      return parser.parse_known_args(argv)
    53  
    54  
    55  def run_command(duration, timeout_signal, command):
    56      """Run a subprocess.  If a timeout elapses, send specified signal.
    57  
    58      :param duration: Timeout in seconds.
    59      :param timeout_signal: Signal to send to the subprocess on timeout.
    60      :param command: Subprocess to run (Popen args).
    61      :return: exit status of the subprocess, 124 if the subprocess was
    62          signalled.
    63      """
    64      if sys.platform == 'win32':
    65          # support CTRL_BREAK
    66          creationflags = subprocess.CREATE_NEW_PROCESS_GROUP
    67      else:
    68          creationflags = 0
    69      proc = subprocess.Popen(command, creationflags=creationflags)
    70      for remaining in chain([None], until_timeout(duration)):
    71          result = proc.poll()
    72          if result is not None:
    73              return result
    74          time.sleep(0.1)
    75      else:
    76          proc.send_signal(timeout_signal)
    77          proc.wait()
    78          return 124
    79  
    80  
    81  def main(args=None):
    82      args, command = parse_args(args)
    83      return run_command(args.duration, signals[args.signal], command)
    84  
    85  
    86  if __name__ == '__main__':
    87      sys.exit(main())