github.com/canonical/ubuntu-image@v0.0.0-20240430122802-2202fe98b290/tests/lib/external/snapd-testing-tools/tools/retry (about)

     1  #!/usr/bin/env python3
     2  from __future__ import print_function, absolute_import, unicode_literals
     3  
     4  import argparse
     5  import itertools
     6  import os
     7  import subprocess
     8  import sys
     9  import time
    10  
    11  
    12  # Define MYPY as False and use it as a conditional for typing import. Despite
    13  # this declaration mypy will really treat MYPY as True when type-checking.
    14  # This is required so that we can import typing on Python 2.x without the
    15  # typing module installed. For more details see:
    16  # https://mypy.readthedocs.io/en/latest/common_issues.html#import-cycles
    17  MYPY = False
    18  if MYPY:
    19      from typing import List, Text
    20  
    21  def envpair(s):
    22      # type: (str) -> str
    23      if not "=" in s:
    24          raise argparse.ArgumentTypeError("environment variables expected format is 'KEY=VAL' got '{}'".format(s))
    25      return s
    26  
    27  def _make_parser():
    28      # type: () -> argparse.ArgumentParser
    29      parser = argparse.ArgumentParser(
    30          description="""
    31  Retry executes COMMAND at most N times, waiting for SECONDS between each
    32  attempt. On failure the exit code from the final attempt is returned.
    33  """
    34      )
    35      parser.add_argument(
    36          "-n",
    37          "--attempts",
    38          metavar="N",
    39          type=int,
    40          default=3,
    41          help="number of attempts (default %(default)s)",
    42      )
    43      parser.add_argument(
    44          "--wait",
    45          metavar="SECONDS",
    46          type=float,
    47          default=1,
    48          help="grace period between attempts (default %(default)ss)",
    49      )
    50      parser.add_argument(
    51          "--env",
    52          type=envpair,
    53          metavar='KEY=VAL',
    54          action='append',
    55          default=[],
    56          help="environment variable to use with format KEY=VALUE (no default)",
    57      )
    58      parser.add_argument(
    59          "--maxmins",
    60          metavar="MINUTES",
    61          type=float,
    62          default=0,
    63          help="number of minutes after which to give up (no default, if set attempts is ignored)",
    64      )
    65      parser.add_argument(
    66          "--quiet",
    67          dest="verbose",
    68          action="store_false",
    69          default=True,
    70          help="refrain from printing any output",
    71      )
    72      parser.add_argument(
    73          "cmd", metavar="COMMAND", nargs="...", help="command to execute"
    74      )
    75      return parser
    76  
    77  
    78  def get_env(env):
    79      # type: (List[str]) -> dict[str,str]
    80      new_env = os.environ.copy()
    81      maxsplit=1  # no keyword support for str.split() in py2
    82      for key, val in [s.split("=", maxsplit) for s in env]:
    83          new_env[key] = val
    84      return new_env
    85  
    86  
    87  def run_cmd(cmd, n, wait, maxmins, verbose, env):
    88      # type: (List[Text], int, float, float, bool, List[str]) -> int
    89      if maxmins != 0:
    90          attempts = itertools.count(1)
    91          t0 = time.time()
    92          after = "{} minutes".format(maxmins)
    93          of_attempts_suffix = ""
    94      else:
    95          attempts = range(1, n + 1)
    96          after = "{} attempts".format(n)
    97          of_attempts_suffix = " of {}".format(n)
    98      retcode = 0
    99      i = 0
   100      new_env = get_env(env)
   101      for i in attempts:
   102          retcode = subprocess.call(cmd, env=new_env)
   103          if retcode == 0:
   104              return 0
   105          if verbose:
   106              print(
   107                  "retry: command {} failed with code {}".format(" ".join(cmd), retcode),
   108                  file=sys.stderr,
   109              )
   110          if maxmins != 0:
   111              elapsed = (time.time()-t0)/60
   112              if elapsed > maxmins:
   113                  break
   114          if i < n or maxmins != 0:
   115              if verbose:
   116                  print(
   117                      "retry: next attempt in {} second(s) (attempt {}{})".format(
   118                          wait, i, of_attempts_suffix
   119                      ),
   120                      file=sys.stderr,
   121                  )
   122              time.sleep(wait)
   123  
   124      if verbose and i > 1:
   125          print(
   126              "retry: command {} keeps failing after {}".format(
   127                  " ".join(cmd), after
   128                  ),
   129                  file=sys.stderr,
   130              )
   131      return retcode
   132  
   133  
   134  def main():
   135      # type: () -> None
   136      parser = _make_parser()
   137      ns = parser.parse_args()
   138      # The command cannot be empty but it is difficult to express in argparse itself.
   139      if len(ns.cmd) == 0:
   140          parser.print_usage()
   141          parser.exit(0)
   142      # Return the last exit code as the exit code of this process.
   143      try:
   144          retcode = run_cmd(ns.cmd, ns.attempts, ns.wait, ns.maxmins, ns.verbose, ns.env)
   145      except OSError as exc:
   146          if ns.verbose:
   147              print(
   148                  "retry: cannot execute command {}: {}".format(" ".join(ns.cmd), exc),
   149                  file=sys.stderr,
   150              )
   151          raise SystemExit(1)
   152      else:
   153          raise SystemExit(retcode)
   154  
   155  
   156  if __name__ == "__main__":
   157      main()