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.
5 """Purpose of this module is to hold common script/commandline functionality.
7 This ranges from optparse, to a basic script wrapper setup (much like
8 what is used for chromite.bin.*).
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
32 CHECKOUT_TYPE_UNKNOWN = 'unknown'
33 CHECKOUT_TYPE_GCLIENT = 'gclient'
34 CHECKOUT_TYPE_REPO = 'repo'
35 CHECKOUT_TYPE_SUBMODULE = 'submodule'
38 CheckoutInfo = collections.namedtuple(
39 'CheckoutInfo', ['type', 'root', 'chrome_src_dir'])
42 def DetermineCheckout(cwd):
43 """Gather information on the checkout we are in.
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.
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
59 gclient_file = os.path.join(path, '.gclient')
60 if os.path.exists(gclient_file):
61 checkout_type = CHECKOUT_TYPE_GCLIENT
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
69 if checkout_type != CHECKOUT_TYPE_UNKNOWN:
72 # Determine the chrome src directory.
74 if checkout_type == CHECKOUT_TYPE_GCLIENT:
75 chrome_src_dir = os.path.join(root, 'src')
76 elif checkout_type == CHECKOUT_TYPE_SUBMODULE:
79 return CheckoutInfo(checkout_type, root, chrome_src_dir)
83 """Calculate the current cache dir.
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.
90 The path to the cache dir.
92 return os.environ.get(
93 constants.SHARED_CACHE_ENVVAR,
94 BaseParser.FindCacheDir(None, None))
97 def AbsolutePath(_option, _opt, value):
98 """Expand paths and make them absolute."""
99 return osutils.ExpandPath(value)
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):]))
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)
114 def ParseDate(value):
115 """Parse date argument into a datetime.date object.
118 value: String representing a single date in "YYYY-MM-DD" format.
121 A datetime.date object.
124 return datetime.datetime.strptime(value, '%Y-%m-%d').date()
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.')
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'):
141 o[2] = os.path.normpath(o[2])
142 return urlparse.urlunparse(o)
144 return NormalizeLocalOrGSPath(value)
147 def OptparseWrapCheck(desc, check_f, _option, opt, value):
148 """Optparse adapter for type checking functionality."""
150 return check_f(value)
152 raise optparse.OptionValueError(
153 'Invalid %s given: --%s=%s' % (desc, opt, value))
158 'path': osutils.ExpandPath,
159 'gs_path': NormalizeGSPath,
160 'local_or_gs_path': NormalizeLocalOrGSPath,
161 'path_or_uri': NormalizeUri,
165 class Option(optparse.Option):
166 """Subclass to implement path evaluation & other useful types."""
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])
175 class FilteringOption(Option):
176 """Subclass that supports Option filtering for FilteringOptionParser"""
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)
184 elif not self.nargs or self.nargs <= 1:
187 parser.AddParsedArg(self, opt, [str(v) for v in value])
190 class BaseParser(object):
191 """Base parser class that includes the logic to add logging controls."""
193 DEFAULT_LOG_LEVELS = ('fatal', 'critical', 'error', 'warning', 'info',
196 DEFAULT_LOG_LEVEL = "info"
199 REPO_CACHE_DIR = '.cache'
200 CHROME_CACHE_DIR = '.cros_cache'
202 def __init__(self, **kwargs):
203 """Initialize this parser instance.
206 logging: Defaults to ALLOW_LOGGING from the class; if given,
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
226 If False, or caching is not given, then no --cache-dir option will be
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)
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',
248 for key in parser_keys:
249 kwarg_dict.pop(key, None)
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',
267 help='Do not use colorized output (or `export NOCOLOR=true`)')
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' .")
276 def SetupLogging(self, opts):
277 value = opts.log_level.upper()
278 logging.getLogger().setLevel(getattr(logging, value))
281 def DoPostParseSetup(self, opts, args):
282 """Method called to handle post opts/args setup.
284 This can be anything from logging setup to positional arg count validation.
287 opts: optparse.Values or argparse.Namespace instance
288 args: position arguments unconsumed from parsing.
291 (opts, args), w/ whatever modification done.
293 if self.logging_enabled:
294 value = self.SetupLogging(opts)
295 if self.debug_enabled:
296 opts.debug = (value == "DEBUG")
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)
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)
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")
318 os.environ[constants.SHARED_CACHE_ENVVAR] = cache_dir
319 logging.debug("Configured cache_dir to %r", cache_dir)
322 def FindCacheDir(cls, _parser, _opts):
323 logging.debug('Cache dir lookup.')
324 checkout = DetermineCheckout(os.getcwd())
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')
333 raise AssertionError('Unexpected type %s' % checkout.type)
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')
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)
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.'
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.'
360 def __init__(self, defaults, *args, **kwargs):
361 cros_build_lib.FrozenAttributesMixin.__init__(self)
362 optparse.Values.__init__(self, defaults, *args, **kwargs)
364 # Used by FilteringParser.
365 self.parsed_args = None
368 class OptionParser(optparse.OptionParser, BaseParser):
369 """Custom parser adding our custom option class in.
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.
378 DEFAULT_OPTION_CLASS = Option
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)
387 def parse_args(self, args=None, values=None):
388 # If no Values object is specified then use our custom OptionValues.
390 values = OptionValues(defaults=self.defaults)
392 opts, remaining = optparse.OptionParser.parse_args(
393 self, args=args, values=values)
394 return self.DoPostParseSetup(opts, remaining)
397 PassedOption = collections.namedtuple(
398 'PassedOption', ['opt_inst', 'opt_str', 'value_str'])
401 class FilteringParser(OptionParser):
402 """Custom option parser for filtering options."""
404 DEFAULT_OPTION_CLASS = FilteringOption
406 def parse_args(self, args=None, values=None):
407 # If no Values object is specified then use our custom OptionValues.
409 values = OptionValues(defaults=self.defaults)
411 values.parsed_args = []
413 return OptionParser.parse_args(self, args=args, values=values)
415 def AddParsedArg(self, opt_inst, opt_str, value_str):
416 """Add a parsed argument with attributes.
419 opt_inst: An instance of a raw optparse.Option object that represents the
421 opt_str: The option string.
422 value_str: A list of string-ified values dentified by OptParse.
424 self.values.parsed_args.append(PassedOption(opt_inst, opt_str, value_str))
427 def FilterArgs(parsed_args, filter_fn):
428 """Filter the argument by passing it through a function.
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.
437 A tuple containing two lists - one of accepted arguments and one of
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)
447 return accepted, removed
450 # pylint: disable=R0901
451 class ArgumentParser(BaseParser, argparse.ArgumentParser):
452 """Custom argument parser for use by chromite.
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.
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)
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)
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)
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)
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()
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]
494 class _ShutDownException(SystemExit):
495 """Exception raised when user hits CTRL+C."""
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)
505 def _DefaultHandler(signum, _frame):
506 # Don't double process sigterms; just trigger shutdown from the first
508 signal.signal(signum, signal.SIG_IGN)
509 raise _ShutDownException(
510 signum, "Received signal %i; shutting down" % (signum,))
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.
518 Note that this function invokes sys.exit on the way out by default.
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.
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)
541 print >> sys.stderr, ("Internal error detected- no main "
542 "functor found in module %r." % (name,))
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.
551 datefmt=constants.LOGGER_DATE_FMT)
553 signal.signal(signal.SIGTERM, _DefaultHandler)
557 ret = target(argv[1:])
558 except _ShutDownException as e:
560 print >> sys.stderr, ("%s: Signaled to shutdown: caught %i signal." %
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.
567 except Exception as e:
569 print >> sys.stderr, ("%s: Unhandled exception:" % (name,))