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