Coverage for /opt/anaconda/lib/python3.7/site-packages/llama/cli.py : 40%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# (c) Stefan Countryman, 2019
3"""
4Command-Line Interface (CLI) primitives to be used by scripts throughout the
5pipeline. Base your scripts off of these classes as much as possible to shorten
6development time and provide a unified "look-and-feel" to the entire library's
7CLI.
8"""
10import re
11import os
12import sys
13from io import StringIO
14from glob import glob
15from copy import deepcopy
16import importlib
17from shlex import quote
18from signal import signal, SIGINT, SIGTERM
19import logging
20import traceback
21import atexit
22from textwrap import wrap, indent
23from pkgutil import iter_modules
24from types import ModuleType, FunctionType
25from typing import Dict, List
26from argparse import (
27 ArgumentParser,
28 RawDescriptionHelpFormatter,
29 Namespace,
30 Action,
31)
32import psutil
33import functools
34from llama.com.slack import alert_maintainers
35from llama.classes import ImmutableDict, EnvVarRegistry
36from llama.utils import (
37 COLOR as COL,
38 setup_logger,
39 tbhighlight,
40 __version__,
41)
43LOGGER = logging.getLogger(__name__)
44LOGLEVELS = ('debug', 'info', 'warning', 'error', 'critical', 'none')
45# stuff for printing out lists of processes
46PROC_HEAD = "PID MEM% RUNNING COMMAND-LINE-OPTS"
47PROC_FMT = "{ind}{PID:<5d} {MEM:<5.2f} {RUNNING:<13} {CMD:<65s}"
48NEXT_LEAF = r'└── '
49CONTINUE_BRANCH = '│ '
50_ERR_ALERT = [False]
53def proc_formatter(pid):
54 """Get a dictionary that can be used to format ``PROC_FMT`` via
55 ``proc_printer`` with information about the process with id ``pid``."""
56 proc = psutil.Process(pid)
57 try:
58 return {
59 'PID': pid,
60 'RUNNING': format(proc.is_running()),
61 'CMD': proc.cmdline(),
62 'NAME': proc.name(),
63 'MEM': proc.memory_percent(),
64 'CHILDREN': [proc_formatter(c.pid) for c in proc.children()],
65 }
66 except psutil.NoSuchProcess:
67 return {
68 'PID': -1,
69 'RUNNING': 'Non-existent!',
70 'CMD': "N/A",
71 'NAME': "Killed.",
72 "MEM": 0.0,
73 "CHILDREN": [],
74 }
77def proc_printer(
78 procs,
79 command_nicknames: Dict[tuple, str] = ImmutableDict({}),
80 indent: list = []
81):
82 """Print a bunch of processes in a nice tree-like way to STDOUT."""
83 for i, proc in enumerate(procs):
84 if not indent:
85 print(PROC_HEAD)
86 if i == len(procs)-1:
87 next_ind = indent[:-1] + [' '*4, NEXT_LEAF]
88 else:
89 next_ind = indent[:-1] + [CONTINUE_BRANCH, NEXT_LEAF]
90 children = proc["CHILDREN"]
91 copy = deepcopy(proc)
92 # need to iterate through first few possible commands since these
93 # can be something like ``python -m``
94 nargs = len(copy['CMD'])
95 for cmdtup, nickname in command_nicknames.items():
96 cmd = list(cmdtup)
97 for i in range(1+nargs-len(cmd)):
98 if cmd == copy['CMD'][i:i+len(cmd)]:
99 copy['CMD'] = (f'[{nickname}] ' +
100 ' '.join(quote(w)
101 for w in copy['CMD'][i+len(cmd)]))
102 break # found it, stop trying this nickname
103 else:
104 continue # no breaks, nickname not found, go to the next
105 break # nickname found, break this loop
106 else: # no breaks, no nicknames found
107 copy['CMD'] = ' '.join(quote(w) for w in copy['CMD'])
108 print(PROC_FMT.format(ind=''.join(indent), **copy))
109 proc_printer(children, indent=next_ind)
112def printprocs(pids, command_nicknames=ImmutableDict({})):
113 """Print a nicely-formatted list of processes and their subprocesses using
114 ``proc_printer``.
116 Parameters
117 ----------
118 pids : list-like,
119 An iterable of process ids.
120 command_nicknames : Dict[tuple, str], optional
121 Specify nicknames for commands as a dictionary mapping from the
122 arguments associated with a command (e.g. ``['llama', 'run']`` to
123 the replacement nicknames for each (each of which will be wrapped in
124 square brackets like ``[llama run]``) when the processes are printed.
125 The ``python`` version printed before these arguments is omitted to
126 save space. This command shortening is intended to highlight
127 and shorten interesting commands.
128 """
129 proc_printer([proc_formatter(p) for p in pids], command_nicknames)
132def register_exit_handler():
133 """Make sure we perform exit actions by calling ``exit`` explicitly on
134 interrupts and SIGTERMs."""
135 def handle_signal(signum, frame):
136 """Log the interrupt/termination attempt and exit cleanly."""
137 try:
138 LOGGER.critical("Received signal %s, cleaning up and exiting NOW.",
139 signum)
140 finally: # pylint: disable=bare-except
141 exit(67)
142 signal(SIGINT, handle_signal)
143 signal(SIGTERM, handle_signal)
146def safe_launch_daemon(lockdir: str, post_fork: FunctionType = None):
147 """Fork this process twice, exiting the parent and grandparent, to put this
148 script into the background and creating a lock directory (atomic on most
149 filesystems) specific to this process, then run a ``post_fork`` function to
150 do any extra initialization or conflict checking (e.g. to check whether
151 this process is conflicting with other processes in a way not accounted for
152 by the initial lock aquisition check). Continues execution in the new
153 grandchild process."""
154 register_exit_handler()
155 spawn_daemon()
156 pidpath = pidfile(lockdir)
157 if post_fork is None:
158 def post_fork():
159 pass
160 # file creation is atomic on most filesystems. check whether a similar
161 # process is running by checking the existence of its lockdir; exit if that
162 # directory exists.
163 try:
164 os.mkdir(lockdir)
166 # register a cleanup function to remove lock directory at exit
167 @atexit.register
168 def cleanup(): # pylint: disable=unused-variable
169 """Remove lock directory if we quit early."""
170 try:
171 os.remove(pidpath)
172 except OSError:
173 LOGGER.warning("Could not remove pidfile %s on exit.", pidpath)
174 try:
175 os.rmdir(lockdir)
176 except Exception as err:
177 LOGGER.error("Could not remove lockdir %s; your next "
178 "invocation of this script will not be able to "
179 "acquire file lock to launch in daemon mode.",
180 lockdir)
181 LOGGER.error(traceback.format_exc())
182 if not isinstance(err, OSError):
183 raise
185 with open(pidpath, 'w') as pidf:
186 pidf.write(str(os.getpid()))
187 post_fork()
188 close_stdout_stderr()
189 except OSError:
190 LOGGER.error("EXITING; Could not acquire lock; is another ``llama "
191 "run`` instance running? Lockdir path: %s", lockdir)
192 exit(1)
195def close_stdout_stderr(outfile=os.devnull):
196 """Redirect stdout and stderr to ``outfile`` at the file descriptor level
197 so that, when the process is forked to the background, no new data is
198 written to stdout/stderr. Since everything should be getting logged anyway,
199 this should not be necessary, but unfortunately we need to use production
200 code that prints to these file descriptors instead of using ``logging``.
202 Parameters
203 ----------
204 outfile : str, optional
205 The path to the logfile that should collect all of the output printed
206 to STDOUT and STDERR. Defaults to ``os.devnull`` (i.e. delete all
207 outputs).
208 """
209 LOGGER.info("CLOSING STDOUT and STDERR now! Redirecting to %s", outfile)
210 sys.stdout.flush()
211 sys.stderr.flush()
212 with open(outfile, 'ab') as redirect:
213 os.dup2(redirect.fileno(), sys.stdout.fileno())
214 os.dup2(redirect.fileno(), sys.stderr.fileno())
217def spawn_daemon():
218 """Do the UNIX double-fork magic, see Stevens' "Advanced
219 Programming in the UNIX Environment" for details (ISBN 0201563177).
220 Taken from https://stackoverflow.com/questions/6011235."""
221 try:
222 pid = os.fork()
223 if pid > 0:
224 # parent process, exit immediately after forking
225 exit()
226 except OSError as err:
227 LOGGER.error("fork #1 failed: %d (%s)", err.errno, err.strerror)
228 exit(1)
230 os.setsid()
232 # do second fork
233 try:
234 pid = os.fork()
235 if pid > 0:
236 # exit from second parent
237 exit()
238 except OSError as err:
239 LOGGER.error("fork #2 failed: %d (%s)", err.errno, err.strerror)
240 exit(1)
242 LOGGER.debug("Put 'llama run' process in background, PID: %s", pid)
245def pidfile(lockdir):
246 """Get the path to the file containing the process ID for the ``llama run``
247 instance running on this run directory with this eventidfilter."""
248 return os.path.join(lockdir, 'pid')
251def running_pids(running_lock_dir: str):
252 """Find all running instances of whatever program is using
253 ``running_lock_dir`` to store its lock directories (each of which should
254 contain a pidfile with the process ID stored in it)."""
255 pids = []
256 for pidfile in glob(os.path.join(running_lock_dir, '*', 'pid')):
257 with open(pidfile) as pfile:
258 pids.append(int(pfile.read()))
259 return pids
262def print_running_procs_action(running_lock_dir: str,
263 command_nicknames: Dict[tuple, str]):
264 """Get a ``PrintRunningProcsAction`` class that can be used in an
265 ``ArgumentParser`` argument as the action to show which processes
266 whose lockfiles are stored in ``running_lock_dir`` are currently
267 running."""
269 nicknames = ImmutableDict(command_nicknames)
271 class PrintRunningProcsAction(Action):
272 """Print running processes whose lockfiles are stored in {} and
273 quit."""
275 def __call__(self, parser, namespace, values, option_string=None):
276 """Only called if flag explicitly specified."""
277 printprocs(running_pids(running_lock_dir), nicknames)
278 exit()
280 __doc__ = __doc__.format(running_lock_dir)
282 return PrintRunningProcsAction
285class PrintEnvAction(Action):
286 """
287 Print environmental variables loaded by ``llama`` and quit.
288 """
290 def __call__(self, parser, namespace, values, option_string=None):
291 """Print and exit."""
292 EnvVarRegistry.print_and_quit()
295class ErrAlertAction(Action):
296 """
297 Indicate that maintainer alerts should be on, letting maintainers know
298 about broken functionality and providing error tracebacks.
299 """
301 def __call__(self, parser, namespace, values, option_string=None):
302 """
303 Mark error alerts as on.
304 """
305 _ERR_ALERT[0] = True
308class CanonicalPathAction(Action):
309 """
310 Canonicalize a collection of paths with ``os.path.realpath``, remove
311 duplicates, and sort the canonicalized paths so that the full list is an
312 unambiguous representation of the specified values.
313 """
315 def __call__(self, parser, namespace, values, option_string=None):
316 """Canonicalize ``values`` whether it is a list or not. Leave alone if
317 ``None``."""
318 if values is None:
319 setattr(namespace, self.dest, None)
320 if isinstance(values, list) or isinstance(values, tuple):
321 setattr(namespace, self.dest,
322 sorted(set([os.path.realpath(p) for p in values])))
323 else:
324 setattr(namespace, self.dest, os.path.realpath(values))
327class CliParser(ArgumentParser):
328 """
329 Extend ``ArgumentParser`` with postprocessing functions that run on the
330 parsed arguments when ``parse_args`` or ``parse_known_args`` are called to
331 adjust their values or raise errors as necessary.
332 """
334 POSTPROCESSORS = tuple()
336 def __init__(self, *args, parents=tuple(), **kwargs):
337 """Run ``ArgumentParser.__init__`` and then update ``POSTPROCESSORS``
338 with any postprocessors passed as ``parents`` in the order in which
339 they appear in ``parents``."""
340 super().__init__(*args, parents=parents, **kwargs)
341 self.POSTPROCESSORS = sum((tuple(p.POSTPROCESSORS)
342 for p in parents
343 if hasattr(p, 'POSTPROCESSORS')),
344 type(self).POSTPROCESSORS)
346 def error(self, message):
347 """Same as ``ArgumentParser.error`` but with a bright red error
348 message."""
349 super().error(COL.RED+message+COL.CLEAR)
351 def print_help(self, file=None):
352 """Print the help string for command line consumption. Same as
353 ``ArgumentParser.print_help``, but cleans up ReST directives in default
354 output of the parser for improved legibility."""
355 file = file or sys.stdout
356 with StringIO() as fakefile:
357 super().print_help(fakefile)
358 fakefile.seek(0)
359 helpstring = fakefile.read()
360 print(helpstring.replace('.. code::\n\n', '').replace('``', '`'),
361 file=file)
363 def parse_known_args(self, args: List[str] = None, namespace=None):
364 """Parse known arguments and apply all functions in ``POSTPROCESSORS``
365 to the returned ``namespace``. Also return unrecognized arguments.
367 Parameters
368 ----------
369 args : List[str], optional
370 The arguments to parse. Will parse from ``sys.argv`` using
371 ``ArgumentParser.parse_args`` if not provided.
372 namespace : optional
373 Namespace to pass to ``ArgumentParser.parse_args``.
375 Returns
376 -------
377 parsed : argparse.Namespace
378 Arguments with ``self.postprocess`` applied.
379 unrecognized : List[str]
380 Unrecognized arguments.
381 """
382 parsed, unknown = super().parse_known_args(args, namespace)
383 LOGGER.debug("Parsed known args %s from %s", parsed, sys.argv)
384 LOGGER.debug("Unrecognized args: %s", unknown)
385 self.postprocess(parsed)
386 return parsed, unknown
388 def postprocess(self, namespace: Namespace): # pylint: disable=no-self-use
389 """A method that acts on the ``argparse.Namespace`` returned by
390 ``ArgumentParser.parse_args`` and returns the same ``namespace`` with
391 any necessary modifications. A good place to raise exceptions or
392 execute actions based on the full combination of parsed arguments.
393 Works by calling ``self.POSTPROCESSORS`` in order (a tuple of functions
394 with the same signature as the unbound ``self.postprocess`` method).
396 Parameters
397 ----------
398 namespace : argparse.Namespace
399 The return value of ``ArgumentParser.parse_args``.
401 Returns
402 -------
403 namespace : argparse.Namespace
404 The input with any necessary transformations applied.
405 """
406 for processor in self.POSTPROCESSORS:
407 processor(self, namespace)
408 return namespace
411def parse_atom(*args, postprocessors=(), **kwargs):
412 """Create a new ``CliParser`` class with no help documentation added
413 and add a single argument to it.
415 Parameters
416 ----------
417 *args
418 Positional arguments to pass to the new parser's ``add_argument``
419 method.
420 postprocessors : list-like
421 A list of functions to set ``POSTPROCESSORS`` to in the returned
422 ``CliParser``
423 **kwargs
424 Keyword arguments to pass to the new parser's ``add_argument`` method.
426 Returns
427 -------
428 parser : ArgumentParser
429 A new parser with a single argument. Pass this to other new
430 ``ArgumentParser`` instances as one of the ``parents``.
431 """
432 parser = CliParser(add_help=False)
433 parser.add_argument(*args, **kwargs)
434 parser.POSTPROCESSORS = postprocessors
435 return parser
438def loglevel_from_cli_count(loglevel):
439 """Get the label, e.g. ``DEBUG``, corresponding to the number of times the
440 user typed ``-v`` at the command line."""
441 loglevel = min(loglevel, 5) # cap out at 5
442 return {
443 None: None,
444 1: 'CRITICAL',
445 2: 'ERROR',
446 3: 'WARNING',
447 4: 'INFO',
448 5: 'DEBUG',
449 }[loglevel]
452def get_postprocess_required_arg(*arg_names):
453 """Return a postprocessor that prints help if the required argument
454 ``arg_names`` are not specified at the command line by checking whether
455 they're set to a value that evaluates as ``False``. Use this if you want to
456 defer checking for a required argument until postprocessing."""
458 def argcheck(self: CliParser, namespace: Namespace):
459 """Checks whether command line arguments were set for {names}."""
460 missing = [n for n in arg_names if not getattr(namespace, n)]
461 if missing:
462 self.error("the following arguments are required: "
463 f"{', '.join(missing)}")
465 argcheck.__doc__ = argcheck.__doc__.format(names=arg_names)
466 return argcheck
469def postprocess_logging(self: CliParser, namespace: Namespace):
470 """Run ``setup_logger(namespace.logfile, loglevel)`` to set up a
471 logger for this script based on user input."""
472 setup_logger(namespace.logfile, namespace.verbosity.upper())
475def get_logging_cli(default_logfile, default_loglevel='error'):
476 """Create a ``CliParser`` that automatically turns on logging to the
477 specified output file (``--logfile``, always at maximum verbosity) as well
478 as the terminal at the specified verbosity level (``--verbosity``) with
479 ``default_logfile`` and ``default_loglevel`` as defaults. ``verbosity``
480 should be ``none`` if no terminal output is to be printed or else one of
481 the ``logging`` log levels in lower case form (see: ``LOGLEVELS).
482 Again, output will be logged to the ``logfile`` at the maximum verbosity
483 (``DEBUG``) to make sure nothing is lost; suppress this behavior by setting
484 ``/dev/null`` as the logfile.
486 Parameters
487 ----------
488 default_logfile : str
489 Path to which the script should log by default.
490 default_loglevel : int or NoneType, optional
491 How verbose to be by default. ``none`` means to print nothing; other
492 values are typical log levels. Must be a value specified in LOGLEVELS.
494 Returns
495 -------
496 parser : CliParser
497 A parser to use as one of the ``parents`` to a new ``CliParser``
498 instance, which will inherit the logging CLI options and automatic
499 logging setup behavior.
501 Raises
502 ------
503 ValueError
504 If ``default_loglevel`` is not a value in ``LOGLEVELS``.
505 """
506 if default_loglevel not in LOGLEVELS:
507 raise ValueError("default_loglevel must be in LOGLEVELS. instead, "
508 f"got: {default_loglevel}")
509 parser = CliParser(add_help=False)
510 parser.POSTPROCESSORS = (postprocess_logging,)
511 log = parser.add_argument_group('logging settings')
512 log.add_argument('-l', '--logfile', default=default_logfile, help=f"""
513 File where logs should be written. By default, all logging produced by
514 ``llama run`` goes to both an archival logfile shared by all instances
515 of the process as well as STDERR. The archival logfile can be
516 overridden with this argument. If you specify ``/dev/null`` or a path
517 that resolves to the same, logfile output will be suppressed
518 automatically. Logs written to the logfile are always at maximum
519 verbosity, i.e. DEBUG. (default: {default_logfile})""")
520 log.add_argument('-v', '--verbosity', default=default_loglevel,
521 choices=LOGLEVELS, help=f"""
522 Set the verbosity level at which to log to STDOUT; the ``--logfile``
523 will ALWAYS receive maximum verbosity logs (unless it is completely
524 supressed by writing to /dev/null). Available choices correspond to
525 logging severity levels from the ``logging`` library, with the addition
526 of ``none`` if you want to completely suppress logging to standard out.
527 (default: {default_loglevel})""")
528 return parser
531def postprocess_dev_mode(self: CliParser, namespace: Namespace):
532 """If we're not running on a clean git repository, quit (unless
533 ``namespace.dev_mode`` is ``True``, indicating that the developer knows the
534 repository is in an unclean state). Intended to help reproducibility and
535 correctness."""
536 if re.match(r'^[0-9]+\.[0-9]+\.[0-9]+$', __version__,
537 flags=re.MULTILINE) or namespace.dev_mode:
538 return
539 self.error("\nERROR: You're using a developmental code version: "
540 f"{__version__}\n You need to use a properly tagged version "
541 "(something like v1.2.3). If this is installed in "
542 "development mode, you'll need to commit your changed and tag "
543 "the last commit with the next sequential versioin number; if "
544 "you installed using a package manager, you'll need to `pip "
545 "install` a properly tagged version. If you are a developer "
546 "and want to run anyway, use ``--dev-mode``/``-D`` to ignore "
547 "this (NOT RECOMMENDED FOR PRODUCTION USE).")
550def postprocess_version(self: CliParser, namespace: Namespace):
551 """If ``namespace.version`` is ``True``, print the LLAMA version and
552 exit."""
553 if namespace.version:
554 print(__version__)
555 exit(0)
558class Parsers:
559 """
560 Use a ``Parsers`` instance to access ``ArgumentParser`` classes that can
561 be passed as a list in any combination to a new ``ArgumentParser`` instance
562 as the ``parents`` keyword argument. This prevents you from having to write
563 the same help documentation repeatedly. You can override any keyword
564 arguments
565 """
567 erralert = parse_atom("--err-alert", action=ErrAlertAction, nargs=0,
568 help="Alert maintainers to unhandled exceptions.")
570 helpenv = parse_atom("--help-env", action=PrintEnvAction, nargs=0, help="""
571 Print a list of environmental variables used by LLAMA, including
572 whether they were loaded from the environment, whether fallback
573 defaults were specified, and their corresponding descriptive/warning
574 messages for when they are not specified. **DO NOT RUN THIS IF RESULTS
575 ARE BEING LOGGED**, since it will include any access credentials
576 defined in the environment.""")
578 outdir = parse_atom("-d", "--outdir", help="""
579 The directory in which to save the generated file. Will be created if
580 it does not exist. If not specified, a default will be used.""")
582 # TODO make this into -f --filename and only use it for the filename
583 # itself. specify the output directory with --outdir.
584 outfile = parse_atom("-o", "--outfile", help="""
585 The name of the output file. Should only be the filename (specify
586 directory with --outdir). If not specified, a default will be used.""")
588 clobber = parse_atom("-c", "--clobber", action="store_true", help="""
589 Delete the file if it already exists. By default, existing files will
590 not be overwritten.""")
592 version = CliParser(add_help=False)
593 version.add_argument('-V', '--version', action='store_true', help="""
594 Print the version number and exit.""")
595 version.POSTPROCESSORS = (postprocess_version,)
597 dev_mode = CliParser(add_help=False, parents=(version,))
598 dev_mode.add_argument('-D', '--dev-mode', action='store_true', help="""
599 If specified, allow the program to run even if the LLAMA version does
600 not conform to semantic version naming. You should not do this during
601 production except in an emergency. If the flag is not specified but
602 local changes to the source code exist, ``llama run`` will complain and
603 quit immediately (the default behavior).""")
604 dev_mode.POSTPROCESSORS = (postprocess_dev_mode,)
607class RecursiveCli:
608 """
609 A recursive command line interface that allows the user to access the
610 ``__main__.py:main()`` functions of submodules using a ``CMD SUBCMD``
611 notation with clever recursive helpstring documentation to enable
612 straightforward subcommand discovery by the user and to avoid cluttering a
613 namespace with hyphen-separated subcommands.
615 Examples
616 --------
617 Using ``RecursiveCli`` to implement ``main`` in ``llama.__main__`` lets you
618 access ``llama.files.__main__:main()`` in a convenient way from the command
619 line. The commands below are equivalent:
621 .. code::
623 python -m llama.files
624 python -m llama files
626 This becomes useful when you realize that you now only need a single
627 script/alias/entry point for your script's submodules. So if you add an
628 entry point or launcher script for ``llama`` to your distribution, you can
629 replace ``python -m llama`` with ``llama``, and you get the ``llama files``
630 subcommand without any extra work. With the addition of a single entry
631 point, the above command becomes:
633 .. code::
635 llama files
637 Better still, this can be applied recursively to every subpackage to access
638 its mixture of package-level CLI commands and submodule CLIs (with no
639 changes to your single point mentioned above). So, for example, the
640 following two can be equivalent:
642 .. code::
644 python -m llama.files.i3
645 llama files i3
647 Most importantly, this automatic feature discovery enables users to
648 heirarchically find commands and features without necessitating changes to
649 higher-level packages docstrings or CLIs. The same code that enables the
650 subcommand syntax shown above allows those subcommands to be listed with a
651 help string, so that the following command will *tell you* that
652 ``llama files`` is a valid command and *summarize* what it does:
654 .. code::
656 llama --help
658 You can then, of course, run ``llama files --help`` to learn details about
659 that module (and see which submodules it offers). This is similar to
660 ``git`` and other programs, but it can be recursed ad-infinitum for rich
661 libraries.
662 """
663 prog: str
664 description: str
665 subcommands: Dict['str', ModuleType]
666 localparser: ArgumentParser
667 preprocessor: FunctionType
669 def __init__(
670 self,
671 prog: str = None,
672 description: str = None,
673 subcommands: Dict['str', ModuleType] = None,
674 localparser: ArgumentParser = None,
675 preprocessor: FunctionType = None
676 ):
677 """
678 Parameters
679 ----------
680 prog : str, optional
681 The program name to use in help documentation. If not specified,
682 the default ``argparse.ArgumentParser`` name will be used. **NB:
683 the default will give potentially inconsistent or incorrect names
684 depending on usage. You should consider overriding it here or by
685 calling ``from_module`` to initialize.**
686 description : str, optional
687 An optional description used to generate this script's ``--help``
688 documentation. You probably want to set this to ``__doc__`` (which
689 should contain an introduction to the package's functionality
690 anyway).
691 subcommands : Dict['str', ModuleType], optional
692 A dictionary mapping the names of subcommands to their associated
693 modules. If you call ``self.main()`` with ``sys.argv`` set to
694 ``['subcmd', 'arg1', 'arg2']``, and ``subcmd`` is a key in
695 ``subcommands``, then ``'subcmd'`` will be removed ``sys.argv``
696 (giving ``sys.argv=['arg1', 'arg2']``) and
697 ``self.subcommands['subcmd'].__main__.main()`` will be called. The
698 first sentence of the ``__doc__`` attribute of each value in
699 ``subcommands`` will be used to set the help documentation for its
700 corresponding key so that a summary of the subcommands is provided
701 in ``self.print_help()``.
702 localparser : ArgumentParser, optional
703 If you have extra local command options for this specific script,
704 you can define them in an ``argparse.ArgumentParser`` instance and
705 provide it here. ``subcommands`` will be added to this instance.
706 **THIS PARSER MUST NOT ACCEPT POSITIONAL ARGUMENTS;** positional
707 arguments are reserved for ``RecursiveCli``.
708 preprocessor : FunctionType, optional
709 If you have a local script functionality that should run (which
710 would usually be encapsulated in this script's ``main`` function),
711 specify it here. This command should accept a single parameter,
712 ``args``. It will be called before any ``subcommands`` are
713 invoked with the fully parsed arguments provided by ``localparser``
714 (if defined) augmented by ``subcommands``; if ``preprocessor`` does
715 not exit, the selected ``subcommand`` will then be called as
716 specified above in the ``subcommands`` description.
718 Raises
719 ------
720 TypeError
721 If ``preprocessor`` is specified but does not accept any arguments.
722 """
723 preprocessor = preprocessor or self.print_help_if_no_cli_args
724 if not preprocessor.__code__.co_varnames:
725 raise TypeError("If specified, ``preprocessor`` must take a "
726 f"positional argument. got: {preprocessor}")
727 self.prog = prog # will be None if not overridden above
728 self.description = description or ""
729 self.subcommands = subcommands or dict()
730 self.localparser = localparser # will be None if not overridden above
731 self.preprocessor = preprocessor
733 @classmethod
734 def from_module(cls, modulename: str, **kwargs):
735 """Autogenerate the ``submodules`` dictionary by finding available CLIs
736 in ``module``.
738 Parameters
739 ----------
740 modulename : str
741 fully-qualified module name in which to search for ``subcommands``
742 to pass to ``__init__``. The keys of ``subcommands`` will simply
743 be the module names of the submodules of ``modulename``. ``prog``
744 will be the ``modulename`` with periods replaced by spaces, e.g.
745 ``'llama.files'`` will turn into ``prog='llama files'``, since this
746 reflects the way the command is used (assuming the top level
747 module, ``llama`` in the given example, implements a
748 ``RecursiveCli`` and is callable from the command line using the
749 top-level module name, again, ``llama`` in this example).
750 **kwargs
751 Remaining arguments (besides ``subcommands`` and ``prog``) will be
752 passed to ``__init__`` along with the subcommands described above.
754 Returns
755 -------
756 cli : RecursiveCli
757 A new instance with ``subcommands`` automatically discovered from
758 ``module``.
760 Raises
761 ------
762 TypeError
763 If ``subcommands`` or ``prog`` is in ``**kwargs`` or if any other
764 arguments not recognized by ``__init__`` are passed.
765 """
766 for conflicting in ['subcommands', 'prog']:
767 if conflicting in kwargs:
768 raise TypeError("Don't specify '{conflicting}' as a keyword "
769 "argument; its value is calculated from "
770 "'modulename'. If you need to manually "
771 "specify it, you can modify ``{conflicting}`` "
772 "after initialization.")
773 subcommands = dict()
774 prog = modulename.replace('.', ' ')
775 locs = importlib.util.find_spec(modulename).submodule_search_locations
776 for _, subname, ispkg in iter_modules(locs):
777 if not ispkg:
778 continue
779 spec = importlib.util.find_spec(f"{modulename}.{subname}")
780 for _, mname, mispkg in iter_modules(
781 spec.submodule_search_locations
782 ):
783 if mname == '__main__' and not mispkg:
784 fullname = f"{modulename}.{subname}.__main__"
785 mpkg = importlib.import_module(fullname)
786 if hasattr(mpkg, 'main'):
787 subcommands[subname] = mpkg
788 return cls(subcommands=subcommands, prog=prog, **kwargs)
790 def get_epilog(self):
791 """Get an ``epilog`` to pass to ``argparse.ArgumentParser`` that
792 contains information on available subcommands."""
793 epi = 'subcommand summaries (pick one, use ``--help`` for more info):\n'
794 for cmd, module in self.subcommands.items():
795 epi += f' {cmd}\n'
796 sentence = (module.__doc__ or "no description.").split('.')[0]+'.'
797 epi += indent(
798 '\n'.join(wrap(sentence.strip(), 72)),
799 ' '*4
800 ) + '\n'
801 return epi
803 def get_parser(self):
804 """Get a command-line argument parser for this CLI.
806 Returns
807 -------
808 parser : CliParser
809 An ``CliParser`` instance (see ``CliParser``, a subclass of
810 ``argparse.ArgumentParser`` implementing post parsing hooks)
811 containing all specified ``subcommands``. If ``self.localparser``
812 was passed at initialization time, then ``parser`` will be
813 initialized therefrom.
814 """
815 parents = [Parsers.version, Parsers.helpenv]
816 if self.localparser:
817 parents.append(self.localparser)
818 parser = CliParser(description=self.description, parents=parents,
819 prog=self.prog, epilog=self.get_epilog(),
820 formatter_class=RawDescriptionHelpFormatter,
821 add_help=False)
822 if self.subcommands:
823 parser.add_argument('-h', '--help', action='store_true', help="""
824 If a SUBCOMMAND is provided, run ``--help`` for that subcommand.
825 If no SUBCOMMAND is provided, print this documentation and
826 exit.""")
827 subcmd = parser.add_argument_group(
828 'subcommands (call one with ``--help`` for details on each)')
829 subcmd.add_argument(
830 "subcommand", nargs="?", choices=tuple(self.subcommands),
831 help="""If a subcommand is provided, ALL remaining
832 arguments will be passed to that command and it will be
833 called.""")
834 subcmd.add_argument(
835 "subargs", nargs="*", help="""If SUBCOMMAND takes its own
836 positional arguments, provide them here. This includes sub-sub
837 commands.""")
838 else:
839 parser.add_argument('-h', '--help', action='store_true', help="""
840 Print this help documentation and exit.""")
841 return parser
843 def main(self):
844 """The ``main`` function that you should run if this module is called
845 as a script. Parses the command line options using
846 ``self.get_parser()``, prints the help documentation if no
847 ``SUBCOMMAND`` is specified at the command line, runs
848 ``self.preprocessor(self.get_parser().parse_args())`` and then, if it
849 completes without exiting, runs a ``SUBCOMMAND`` if specified.
850 """
851 parser = self.get_parser()
852 args, _unrecognized_args = parser.parse_known_args()
853 subcommand = getattr(args, 'subcommand', None)
854 if subcommand is None:
855 if args.help:
856 parser.print_help()
857 exit()
858 self.preprocessor(args) # calls help if preprocessor not specified
859 if subcommand is not None:
860 sys.argv.remove(subcommand)
861 self.subcommands[subcommand].main()
863 # This function DOES use self in its class definition; spurious pylint
864 # pylint: disable=no-self-use
865 def print_help_if_no_cli_args(self, _args):
866 """If no command-line arguments are parsed, prints the help
867 documentation and exits. This is the default ``preprocessor`` (since,
868 in the absence of another preprocessor, a complete absence of CLI
869 arguments results in a no-op likely indicating a lack of user
870 understanding).
872 Uses an ``CliParser`` to parse arguments (since it knows how
873 to deal with variations in executable names, e.g. ``python -m
874 llama`` vs. ``llama/__main__.py`` vs. ``python3 llama/__main__.py``)
875 and, if no arguments whatsoever are discernable, runs
876 ``self.get_parser().print_help()`` and quits.
877 """
879 class Quitter(ArgumentParser):
880 """Throwaway class to make sure we parse arguments the exact way
881 that the parser from ``get_parser`` would do it."""
882 def error(_self, message): # pylint: disable=no-self-argument
883 self.get_parser().print_help()
884 exit()
886 parser = Quitter(add_help=False)
887 parser.add_argument("args", nargs="+")
888 parser.parse_known_args()
891def traceback_alert_maintainers(func: FunctionType, err: Exception, tb: str,
892 self, *args, **kwargs):
893 """An ``error_callback`` function for ``log_exceptions_and_recover``.
894 Runs ``alert_maintainers`` with the current traceback and logs the stack
895 trace. Does not send out alerts if ``--err-alert`` is not set at the
896 command line.
898 Parameters
899 ----------
900 func : FunctionType
901 The function that caused the error.
902 err : Exception
903 The exception that was caught.
904 tb : str
905 The traceback to send as a message.
906 self : object or None
907 If ``func`` is a method, this will be the ``__self__`` value bound to
908 it, otherwise ``None``.
909 *args
910 The positional arguments passed to ``func`` that caused ``err``.
911 **kwargs
912 The keyword arguments passed to ``func`` that caused ``err``.
913 """
914 if _ERR_ALERT[0]:
915 LOGGER.error("Attempting to `alert_maintainers` with traceback...")
916 res = alert_maintainers(
917 f"ERROR: {err}. Traceback:\n\n```\n{tb}\n```\n`args`:\n"
918 f"{args}\n`kwargs`:\n{kwargs}\n`self`:\n{self}\n",
919 func.__name__,
920 recover=True,
921 )
922 # traceback will be printed if failure happens while sending alert
923 if res.get('ok', False):
924 LOGGER.debug("Error Alert sent.")
925 else:
926 LOGGER.error("Error Alert failed to send.")
927 else:
928 LOGGER.debug("Error alert off.")
931def log_exceptions_and_recover(callbacks=(traceback_alert_maintainers,)):
932 """Decorator to run ``func`` with no arguments. Log stack trace and run
933 callbacks to clean up (default: send the traceback to maintainers) when any
934 ``Exception`` is raised, returning the value of the wrapped function or
935 else the exception that was caught (note that this will break functionality
936 of any function that is supposed to return an ``Exception``, and that you
937 should only apply this sort of behavior in a command line script that needs
938 to recover from all exceptions). Error tracebacks are syntax-highlighted
939 (for 256-color terminals) for readability.
941 Optionally provide an iterable of callbacks to run to override the default,
942 ``traceback_alert_maintainers``. Callbacks passed in this list must have
943 the same signature as that function. Use this to perform other cleanup
944 tasks or to avoid sending an alert on error.
946 Signature below is for the returned decorator.
948 Parameters
949 ----------
950 func : function
951 A function that is supposed to recover from **all** possible
952 exceptions. Exceptions with tracebacks will be logged and sent to
953 maintainers using ``alert_maintainers``. You should only wrap functions
954 that **cannot** be allowed to crash, e.g. the ``main`` function of a
955 long-running CLI script.
957 Returns
958 -------
959 func : function
960 The wrapped function. Has the same inputs and return values as the
961 original function *unless* the original function raises an
962 ``Exception`` while executing. In this case, the return value will be
963 the exception instance that was raised. Note that you probably don't
964 care about this value in normal use patterns, and also note that you
965 should therefore *not* wrap a function that would ever nominally return
966 an exception instance since there will no longer be any way to
967 distinguish an expected return value from a caught exception.
968 """
970 def decorator(func: FunctionType):
972 @functools.wraps(func)
973 def wrapper(*args, **kwargs):
974 try:
975 return func(*args, **kwargs)
976 # We always need this to be running, so log any exceptions, send an
977 # alert email, and keep moving.
978 except Exception as err:
979 LOGGER.error("Exception while running %s.%s",
980 func.__module__, func.__name__)
981 LOGGER.error("Printing stack trace.")
982 trace = traceback.format_exc()
983 LOGGER.error(tbhighlight(trace))
984 for callback in callbacks:
985 LOGGER.info("Running error callback: %s",
986 callback.__name__)
987 try:
988 callback(func, err, trace,
989 getattr(func, '__self__', None), *args,
990 **kwargs)
991 except Exception:
992 LOGGER.error("Error while running callback, trace:")
993 LOGGER.error(tbhighlight(traceback.format_exc()))
994 LOGGER.error("Recovering, continuing callbacks.")
995 LOGGER.error("Recovering, returning to normal operation.")
996 return err
997 return wrapper
999 return decorator