src/third_party/chromite/lib/commandline.py

   1 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Purpose of this module is to hold common script/commandline functionality.
   6
   7 This ranges from optparse, to a basic script wrapper setup (much like
   8 what is used for chromite.bin.*).
   9 """
  10
  11 import argparse
  12 import collections
  13 import datetime
  14 import functools
  15 import logging
  16 import os
  17 import optparse
  18 import signal
  19 import sys
  20 import tempfile
  21 import urlparse
  22
  23 # TODO(build): sort the cbuildbot.constants/lib.constants issue;
  24 # lib shouldn't have to import from buildbot like this.
  25 from chromite.cbuildbot import constants
  26 from chromite.lib import cros_build_lib
  27 from chromite.lib import git
  28 from chromite.lib import gs
  29 from chromite.lib import osutils
  30
  31
  32 CHECKOUT_TYPE_UNKNOWN = 'unknown'
  33 CHECKOUT_TYPE_GCLIENT = 'gclient'
  34 CHECKOUT_TYPE_REPO = 'repo'
  35 CHECKOUT_TYPE_SUBMODULE = 'submodule'
  36
  37
  38 CheckoutInfo = collections.namedtuple(
  39     'CheckoutInfo', ['type', 'root', 'chrome_src_dir'])
  40
  41
  42 class ChrootRequiredError(Exception):
  43   """Raised when a command must be run in the chroot
  44
  45   This exception is intended to be caught by code which will restart execution
  46   in the chroot. If none of the arguments passed to the script need to be
  47   adjusted when that happens, it can be constructed with no parameters. If
  48   something does need to be adjusted, for instance an argument that's a path,
  49   the command can construct a custom command line and pass it into this
  50   exception which will be used instead.
  51
  52   When customizing the command line, argv[0] will have to be fixed up manually
  53   like any other element of argv.
  54   """
  55
  56   def __init__(self, new_argv=None, *args, **kwargs):
  57     Exception.__init__(self, *args, **kwargs)
  58     if new_argv is None:
  59       new_argv = sys.argv[:]
  60       new_argv = [git.ReinterpretPathForChroot(new_argv[0])] + new_argv[1:]
  61
  62     self.new_argv = new_argv
  63
  64
  65 def DetermineCheckout(cwd):
  66   """Gather information on the checkout we are in.
  67
  68   Returns:
  69     A CheckoutInfo object with these attributes:
  70       type: The type of checkout.  Valid values are CHECKOUT_TYPE_*.
  71       root: The root of the checkout.
  72       chrome_src_dir: If the checkout is a Chrome checkout, the path to the
  73         Chrome src/ directory.
  74   """
  75   checkout_type = CHECKOUT_TYPE_UNKNOWN
  76   root, path = None, None
  77   for path in osutils.IteratePathParents(cwd):
  78     repo_dir = os.path.join(path, '.repo')
  79     if os.path.isdir(repo_dir):
  80       checkout_type = CHECKOUT_TYPE_REPO
  81       break
  82     gclient_file = os.path.join(path, '.gclient')
  83     if os.path.exists(gclient_file):
  84       checkout_type = CHECKOUT_TYPE_GCLIENT
  85       break
  86     submodule_git = os.path.join(path, '.git')
  87     if (os.path.isdir(submodule_git) and
  88         git.IsSubmoduleCheckoutRoot(cwd, 'origin', constants.CHROMIUM_GOB_URL)):
  89       checkout_type = CHECKOUT_TYPE_SUBMODULE
  90       break
  91
  92   if checkout_type != CHECKOUT_TYPE_UNKNOWN:
  93     root = path
  94
  95   # Determine the chrome src directory.
  96   chrome_src_dir = None
  97   if checkout_type == CHECKOUT_TYPE_GCLIENT:
  98     chrome_src_dir = os.path.join(root, 'src')
  99   elif checkout_type == CHECKOUT_TYPE_SUBMODULE:
 100     chrome_src_dir = root
 101
 102   return CheckoutInfo(checkout_type, root, chrome_src_dir)
 103
 104
 105 def GetCacheDir():
 106   """Calculate the current cache dir.
 107
 108   Users can configure the cache dir using the --cache-dir argument and it is
 109   shared between cbuildbot and all child processes. If no cache dir is
 110   specified, FindCacheDir finds an alternative location to store the cache.
 111
 112   Returns:
 113     The path to the cache dir.
 114   """
 115   return os.environ.get(
 116       constants.SHARED_CACHE_ENVVAR,
 117       BaseParser.FindCacheDir(None, None))
 118
 119
 120 def AbsolutePath(_option, _opt, value):
 121   """Expand paths and make them absolute."""
 122   return osutils.ExpandPath(value)
 123
 124
 125 def NormalizeGSPath(value):
 126   """Normalize GS paths."""
 127   url = gs.CanonicalizeURL(value, strict=True)
 128   return '%s%s' % (gs.BASE_GS_URL, os.path.normpath(url[len(gs.BASE_GS_URL):]))
 129
 130
 131 def NormalizeLocalOrGSPath(value):
 132   """Normalize a local or GS path."""
 133   ptype = 'gs_path' if value.startswith(gs.BASE_GS_URL) else 'path'
 134   return VALID_TYPES[ptype](value)
 135
 136
 137 def ParseDate(value):
 138   """Parse date argument into a datetime.date object.
 139
 140   Args:
 141     value: String representing a single date in "YYYY-MM-DD" format.
 142
 143   Returns:
 144     A datetime.date object.
 145   """
 146   try:
 147     return datetime.datetime.strptime(value, '%Y-%m-%d').date()
 148   except ValueError:
 149     # Give a helpful error message about the format expected.  Putting this
 150     # message in the exception is useless because argparse ignores the
 151     # exception message and just says the value is invalid.
 152     cros_build_lib.Error('Date is expected to be in format YYYY-MM-DD.')
 153     raise
 154
 155
 156 def NormalizeUri(value):
 157   """Normalize a local path or URI."""
 158   # Pylint is confused about result of urlparse.
 159   # pylint: disable=E1101
 160   o = urlparse.urlparse(value)
 161   if o.scheme == 'file':
 162     # Trim off the file:// prefix.
 163     return VALID_TYPES['path'](value[7:])
 164   elif o.scheme not in ('', 'gs'):
 165     o = list(o)
 166     o[2] = os.path.normpath(o[2])
 167     return urlparse.urlunparse(o)
 168   else:
 169     return NormalizeLocalOrGSPath(value)
 170
 171
 172 def OptparseWrapCheck(desc, check_f, _option, opt, value):
 173   """Optparse adapter for type checking functionality."""
 174   try:
 175     return check_f(value)
 176   except ValueError:
 177     raise optparse.OptionValueError(
 178         'Invalid %s given: --%s=%s' % (desc, opt, value))
 179
 180
 181 VALID_TYPES = {
 182     'date': ParseDate,
 183     'path': osutils.ExpandPath,
 184     'gs_path': NormalizeGSPath,
 185     'local_or_gs_path': NormalizeLocalOrGSPath,
 186     'path_or_uri': NormalizeUri,
 187 }
 188
 189
 190 class Option(optparse.Option):
 191   """Subclass to implement path evaluation & other useful types."""
 192
 193   _EXTRA_TYPES = ("path", "gs_path")
 194   TYPES = optparse.Option.TYPES + _EXTRA_TYPES
 195   TYPE_CHECKER = optparse.Option.TYPE_CHECKER.copy()
 196   for t in _EXTRA_TYPES:
 197     TYPE_CHECKER[t] = functools.partial(OptparseWrapCheck, t, VALID_TYPES[t])
 198
 199
 200 class FilteringOption(Option):
 201   """Subclass that supports Option filtering for FilteringOptionParser"""
 202
 203   def take_action(self, action, dest, opt, value, values, parser):
 204     if action in FilteringOption.ACTIONS:
 205       Option.take_action(self, action, dest, opt, value, values, parser)
 206
 207     if value is None:
 208       value = []
 209     elif not self.nargs or self.nargs <= 1:
 210       value = [value]
 211
 212     parser.AddParsedArg(self, opt, [str(v) for v in value])
 213
 214
 215 class BaseParser(object):
 216   """Base parser class that includes the logic to add logging controls."""
 217
 218   DEFAULT_LOG_LEVELS = ('fatal', 'critical', 'error', 'warning', 'info',
 219                         'debug')
 220
 221   DEFAULT_LOG_LEVEL = "info"
 222   ALLOW_LOGGING = True
 223
 224   REPO_CACHE_DIR = '.cache'
 225   CHROME_CACHE_DIR = '.cros_cache'
 226
 227   def __init__(self, **kwargs):
 228     """Initialize this parser instance.
 229
 230     kwargs:
 231       logging: Defaults to ALLOW_LOGGING from the class; if given,
 232         add --log-level.
 233       default_log_level: If logging is enabled, override the default logging
 234         level.  Defaults to the class's DEFAULT_LOG_LEVEL value.
 235       log_levels: If logging is enabled, this overrides the enumeration of
 236         allowed logging levels.  If not given, defaults to the classes
 237         DEFAULT_LOG_LEVELS value.
 238       manual_debug: If logging is enabled and this is True, suppress addition
 239         of a --debug alias.  This option defaults to True unless 'debug' has
 240         been exempted from the allowed logging level targets.
 241       caching: If given, must be either a callable that discerns the cache
 242         location if it wasn't specified (the prototype must be akin to
 243         lambda parser, values:calculated_cache_dir_path; it may return None to
 244         indicate that it handles setting the value on its own later in the
 245         parsing including setting the env), or True; if True, the
 246         machinery defaults to invoking the class's FindCacheDir method
 247         (which can be overridden).  FindCacheDir $CROS_CACHEDIR, falling
 248         back to $REPO/.cache, finally falling back to $TMP.
 249         Note that the cache_dir is not created, just discerned where it
 250         should live.
 251         If False, or caching is not given, then no --cache-dir option will be
 252         added.
 253     """
 254     self.debug_enabled = False
 255     self.caching_group = None
 256     self.debug_group = None
 257     self.default_log_level = None
 258     self.log_levels = None
 259     self.logging_enabled = kwargs.get('logging', self.ALLOW_LOGGING)
 260     self.default_log_level = kwargs.get('default_log_level',
 261                                         self.DEFAULT_LOG_LEVEL)
 262     self.log_levels = tuple(x.lower() for x in
 263                             kwargs.get('log_levels', self.DEFAULT_LOG_LEVELS))
 264     self.debug_enabled = (not kwargs.get('manual_debug', False)
 265                           and 'debug' in self.log_levels)
 266     self.caching = kwargs.get('caching', False)
 267
 268   @staticmethod
 269   def PopUsedArgs(kwarg_dict):
 270     """Removes keys used by the base parser from the kwarg namespace."""
 271     parser_keys = ['logging', 'default_log_level', 'log_levels', 'manual_debug',
 272                    'caching']
 273     for key in parser_keys:
 274       kwarg_dict.pop(key, None)
 275
 276   def SetupOptions(self):
 277     """Sets up special chromite options for an OptionParser."""
 278     if self.logging_enabled:
 279       self.debug_group = self.add_option_group("Debug options")
 280       self.add_option_to_group(
 281           self.debug_group, "--log-level", choices=self.log_levels,
 282           default=self.default_log_level,
 283           help="Set logging level to report at.")
 284       if self.debug_enabled:
 285         self.add_option_to_group(
 286           self.debug_group, "--debug", action="store_const", const="debug",
 287           dest="log_level", help="Alias for `--log-level=debug`. "
 288           "Useful for debugging bugs/failures.")
 289       self.add_option_to_group(
 290         self.debug_group, '--nocolor', action='store_false', dest='color',
 291         default=None,
 292         help='Do not use colorized output (or `export NOCOLOR=true`)')
 293
 294     if self.caching:
 295       self.caching_group = self.add_option_group("Caching Options")
 296       self.add_option_to_group(
 297           self.caching_group, "--cache-dir", default=None, type='path',
 298           help="Override the calculated chromeos cache directory; "
 299           "typically defaults to '$REPO/.cache' .")
 300
 301   def SetupLogging(self, opts):
 302     value = opts.log_level.upper()
 303     logging.getLogger().setLevel(getattr(logging, value))
 304     return value
 305
 306   def DoPostParseSetup(self, opts, args):
 307     """Method called to handle post opts/args setup.
 308
 309     This can be anything from logging setup to positional arg count validation.
 310
 311     Args:
 312       opts: optparse.Values or argparse.Namespace instance
 313       args: position arguments unconsumed from parsing.
 314
 315     Returns:
 316       (opts, args), w/ whatever modification done.
 317     """
 318     if self.logging_enabled:
 319       value = self.SetupLogging(opts)
 320       if self.debug_enabled:
 321         opts.debug = (value == "DEBUG")
 322
 323     if self.caching:
 324       path = os.environ.get(constants.SHARED_CACHE_ENVVAR)
 325       if path is not None and opts.cache_dir is None:
 326         opts.cache_dir = os.path.abspath(path)
 327
 328       opts.cache_dir_specified = opts.cache_dir is not None
 329       if not opts.cache_dir_specified:
 330         func = self.FindCacheDir if not callable(self.caching) else self.caching
 331         opts.cache_dir = func(self, opts)
 332       if opts.cache_dir is not None:
 333         self.ConfigureCacheDir(opts.cache_dir)
 334
 335     return opts, args
 336
 337   @staticmethod
 338   def ConfigureCacheDir(cache_dir):
 339     if cache_dir is None:
 340       os.environ.pop(constants.SHARED_CACHE_ENVVAR, None)
 341       logging.debug("Removed cache_dir setting")
 342     else:
 343       os.environ[constants.SHARED_CACHE_ENVVAR] = cache_dir
 344       logging.debug("Configured cache_dir to %r", cache_dir)
 345
 346   @classmethod
 347   def FindCacheDir(cls, _parser, _opts):
 348     logging.debug('Cache dir lookup.')
 349     checkout = DetermineCheckout(os.getcwd())
 350     path = None
 351     if checkout.type == CHECKOUT_TYPE_REPO:
 352       path = os.path.join(checkout.root, cls.REPO_CACHE_DIR)
 353     elif checkout.type in (CHECKOUT_TYPE_GCLIENT, CHECKOUT_TYPE_SUBMODULE):
 354       path = os.path.join(checkout.root, cls.CHROME_CACHE_DIR)
 355     elif checkout.type == CHECKOUT_TYPE_UNKNOWN:
 356       path = os.path.join(tempfile.gettempdir(), 'chromeos-cache')
 357     else:
 358       raise AssertionError('Unexpected type %s' % checkout.type)
 359
 360     return path
 361
 362   def add_option_group(self, *args, **kwargs):
 363     """Returns a new option group see optparse.OptionParser.add_option_group."""
 364     raise NotImplementedError('Subclass must override this method')
 365
 366   @staticmethod
 367   def add_option_to_group(group, *args, **kwargs):
 368     """Adds the given option defined by args and kwargs to group."""
 369     group.add_option(*args, **kwargs)
 370
 371
 372 class ArgumentNamespace(argparse.Namespace):
 373   """Class to mimic argparse.Namespace with value freezing support."""
 374   __metaclass__ = cros_build_lib.FrozenAttributesClass
 375   _FROZEN_ERR_MSG = 'Option values are frozen, cannot alter %s.'
 376
 377
 378 # Note that because optparse.Values is not a new-style class this class
 379 # must use the mixin FrozenAttributesMixin rather than the metaclass
 380 # FrozenAttributesClass.
 381 class OptionValues(cros_build_lib.FrozenAttributesMixin, optparse.Values):
 382   """Class to mimic optparse.Values with value freezing support."""
 383   _FROZEN_ERR_MSG = 'Option values are frozen, cannot alter %s.'
 384
 385   def __init__(self, defaults, *args, **kwargs):
 386     cros_build_lib.FrozenAttributesMixin.__init__(self)
 387     optparse.Values.__init__(self, defaults, *args, **kwargs)
 388
 389     # Used by FilteringParser.
 390     self.parsed_args = None
 391
 392
 393 class OptionParser(optparse.OptionParser, BaseParser):
 394   """Custom parser adding our custom option class in.
 395
 396   Aside from adding a couple of types (path for absolute paths,
 397   gs_path for google storage urls, and log_level for logging level control),
 398   this additionally exposes logging control by default; if undesired,
 399   either derive from this class setting ALLOW_LOGGING to False, or
 400   pass in logging=False to the constructor.
 401   """
 402
 403   DEFAULT_OPTION_CLASS = Option
 404
 405   def __init__(self, usage=None, **kwargs):
 406     BaseParser.__init__(self, **kwargs)
 407     self.PopUsedArgs(kwargs)
 408     kwargs.setdefault("option_class", self.DEFAULT_OPTION_CLASS)
 409     optparse.OptionParser.__init__(self, usage=usage, **kwargs)
 410     self.SetupOptions()
 411
 412   def parse_args(self, args=None, values=None):
 413     # If no Values object is specified then use our custom OptionValues.
 414     if values is None:
 415       values = OptionValues(defaults=self.defaults)
 416
 417     opts, remaining = optparse.OptionParser.parse_args(
 418         self, args=args, values=values)
 419     return self.DoPostParseSetup(opts, remaining)
 420
 421
 422 PassedOption = collections.namedtuple(
 423         'PassedOption', ['opt_inst', 'opt_str', 'value_str'])
 424
 425
 426 class FilteringParser(OptionParser):
 427   """Custom option parser for filtering options."""
 428
 429   DEFAULT_OPTION_CLASS = FilteringOption
 430
 431   def parse_args(self, args=None, values=None):
 432     # If no Values object is specified then use our custom OptionValues.
 433     if values is None:
 434       values = OptionValues(defaults=self.defaults)
 435
 436     values.parsed_args = []
 437
 438     return OptionParser.parse_args(self, args=args, values=values)
 439
 440   def AddParsedArg(self, opt_inst, opt_str, value_str):
 441     """Add a parsed argument with attributes.
 442
 443     Args:
 444       opt_inst: An instance of a raw optparse.Option object that represents the
 445                 option.
 446       opt_str: The option string.
 447       value_str: A list of string-ified values dentified by OptParse.
 448     """
 449     self.values.parsed_args.append(PassedOption(opt_inst, opt_str, value_str))
 450
 451   @staticmethod
 452   def FilterArgs(parsed_args, filter_fn):
 453     """Filter the argument by passing it through a function.
 454
 455     Args:
 456       parsed_args: The list of parsed argument namedtuples to filter.  Tuples
 457         are of the form (opt_inst, opt_str, value_str).
 458       filter_fn: A function with signature f(PassedOption), and returns True if
 459         the argument is to be passed through.  False if not.
 460
 461     Returns:
 462       A tuple containing two lists - one of accepted arguments and one of
 463       removed arguments.
 464     """
 465     removed = []
 466     accepted = []
 467     for arg in parsed_args:
 468       target = accepted if filter_fn(arg) else removed
 469       target.append(arg.opt_str)
 470       target.extend(arg.value_str)
 471
 472     return accepted, removed
 473
 474
 475 # pylint: disable=R0901
 476 class ArgumentParser(BaseParser, argparse.ArgumentParser):
 477   """Custom argument parser for use by chromite.
 478
 479   This class additionally exposes logging control by default; if undesired,
 480   either derive from this class setting ALLOW_LOGGING to False, or
 481   pass in logging=False to the constructor.
 482   """
 483   # pylint: disable=W0231
 484   def __init__(self, usage=None, **kwargs):
 485     kwargs.setdefault('formatter_class', argparse.RawDescriptionHelpFormatter)
 486     BaseParser.__init__(self, **kwargs)
 487     self.PopUsedArgs(kwargs)
 488     argparse.ArgumentParser.__init__(self, usage=usage, **kwargs)
 489     self._SetupTypes()
 490     self.SetupOptions()
 491
 492   def _SetupTypes(self):
 493     """Register types with ArgumentParser."""
 494     for t, check_f in VALID_TYPES.iteritems():
 495       self.register('type', t, check_f)
 496
 497   def add_option_group(self, *args, **kwargs):
 498     """Return an argument group rather than an option group."""
 499     return self.add_argument_group(*args, **kwargs)
 500
 501   @staticmethod
 502   def add_option_to_group(group, *args, **kwargs):
 503     """Adds an argument rather than an option to the given group."""
 504     return group.add_argument(*args, **kwargs)
 505
 506   def parse_args(self, args=None, namespace=None):
 507     """Translates OptionParser call to equivalent ArgumentParser call."""
 508     # If no Namespace object is specified then use our custom ArgumentNamespace.
 509     if namespace is None:
 510       namespace = ArgumentNamespace()
 511
 512     # Unlike OptionParser, ArgParser works only with a single namespace and no
 513     # args. Re-use BaseParser DoPostParseSetup but only take the namespace.
 514     namespace = argparse.ArgumentParser.parse_args(
 515         self, args=args, namespace=namespace)
 516     return self.DoPostParseSetup(namespace, None)[0]
 517
 518
 519 class _ShutDownException(SystemExit):
 520   """Exception raised when user hits CTRL+C."""
 521
 522   def __init__(self, sig_num, message):
 523     self.signal = sig_num
 524     # Setup a usage message primarily for any code that may intercept it
 525     # while this exception is crashing back up the stack to us.
 526     SystemExit.__init__(self, message)
 527     self.args = (sig_num, message)
 528
 529
 530 def _DefaultHandler(signum, _frame):
 531   # Don't double process sigterms; just trigger shutdown from the first
 532   # exception.
 533   signal.signal(signum, signal.SIG_IGN)
 534   raise _ShutDownException(
 535       signum, "Received signal %i; shutting down" % (signum,))
 536
 537
 538 def _RestartInChroot(argv):
 539   """Rerun the current command inside the chroot"""
 540   return cros_build_lib.RunCommand(argv, enter_chroot=True, error_code_ok=True,
 541                                    cwd=constants.SOURCE_ROOT).returncode
 542
 543
 544 def ScriptWrapperMain(find_target_func, argv=None,
 545                       log_level=logging.DEBUG,
 546                       log_format=constants.LOGGER_FMT):
 547   """Function usable for chromite.script.* style wrapping.
 548
 549   Note that this function invokes sys.exit on the way out by default.
 550
 551   Args:
 552     find_target_func: a function, which, when given the absolute
 553       pathway the script was invoked via (for example,
 554       /home/ferringb/cros/trunk/chromite/bin/cros_sdk; note that any
 555       trailing .py from the path name will be removed),
 556       will return the main function to invoke (that functor will take
 557       a single arg- a list of arguments, and shall return either None
 558       or an integer, to indicate the exit code).
 559     argv: sys.argv, or an equivalent tuple for testing.  If nothing is
 560       given, sys.argv is defaulted to.
 561     log_level: Default logging level to start at.
 562     log_format: Default logging format to use.
 563   """
 564   if argv is None:
 565     argv = sys.argv[:]
 566   target = os.path.abspath(argv[0])
 567   name = os.path.basename(target)
 568   if target.endswith('.py'):
 569     target = os.path.splitext(target)[0]
 570   target = find_target_func(target)
 571   if target is None:
 572     print >> sys.stderr, ("Internal error detected- no main "
 573                           "functor found in module %r." % (name,))
 574     sys.exit(100)
 575
 576   # Set up basic logging information for all modules that use logging.
 577   # Note a script target may setup default logging in its module namespace
 578   # which will take precedence over this.
 579   logging.basicConfig(
 580       level=log_level,
 581       format=log_format,
 582       datefmt=constants.LOGGER_DATE_FMT)
 583
 584   signal.signal(signal.SIGTERM, _DefaultHandler)
 585
 586   ret = 1
 587   try:
 588     ret = target(argv[1:])
 589   except _ShutDownException as e:
 590     sys.stdout.flush()
 591     print >> sys.stderr, ("%s: Signaled to shutdown: caught %i signal." %
 592                           (name, e.signal,))
 593     sys.stderr.flush()
 594   except SystemExit as e:
 595     # Right now, let this crash through- longer term, we'll update the scripts
 596     # in question to not use sys.exit, and make this into a flagged error.
 597     raise
 598   except ChrootRequiredError as e:
 599     ret = _RestartInChroot(e.new_argv)
 600   except Exception as e:
 601     sys.stdout.flush()
 602     print >> sys.stderr, ("%s: Unhandled exception:" % (name,))
 603     sys.stderr.flush()
 604     raise
 605   finally:
 606     logging.shutdown()
 607
 608   if ret is None:
 609     ret = 0
 610   sys.exit(ret)