xref: /aosp_15_r20/external/crosvm/tools/impl/cli.py (revision bb4ee6a4ae7042d18b07a98463b9c8b875e44b39)
1#!/usr/bin/env python3
2# Copyright 2023 The ChromiumOS Authors
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""
7Provides a framework for command line interfaces based on argh.
8
9It automatically adds common arguments, such as -v, -vv and --color to provide consistent
10behavior.
11"""
12
13import argparse
14import sys
15import traceback
16from typing import (
17    Any,
18    Callable,
19    Optional,
20)
21
22from .util import (
23    add_common_args,
24    parse_common_args,
25    print_timing_info,
26    record_time,
27    verbose,
28    ensure_packages_exist,
29)
30
31ensure_packages_exist("argh")
32import argh  # type: ignore
33
34# Hack: argh does not support type annotations. This prevents type errors.
35argh: Any  # type: ignore
36
37
38def run_main(main_fn: Callable[..., Any], usage: Optional[str] = None):
39    run_commands(default_fn=main_fn, usage=usage)
40
41
42def run_commands(
43    *functions: Callable[..., Any],
44    default_fn: Optional[Callable[..., Any]] = None,
45    usage: Optional[str] = None,
46):
47    """
48    Allow the user to call the provided functions with command line arguments translated to
49    function arguments via argh: https://pythonhosted.org/argh
50    """
51    exit_code = 0
52    try:
53        parser = argparse.ArgumentParser(
54            description=usage,
55            # Docstrings are used as the description in argparse, preserve their formatting.
56            formatter_class=argparse.RawDescriptionHelpFormatter,
57            # Do not allow implied abbreviations. Abbreviations should be manually specified.
58            allow_abbrev=False,
59        )
60        add_common_args(parser)
61
62        # Add provided commands to parser. Do not use sub-commands if we just got one function.
63        if functions:
64            argh.add_commands(parser, functions)  # type: ignore
65        if default_fn:
66            argh.set_default_command(parser, default_fn)  # type: ignore
67
68        with record_time("Total Time"):
69            # Call main method
70            argh.dispatch(parser)  # type: ignore
71
72    except Exception as e:
73        if verbose():
74            traceback.print_exc()
75        else:
76            print(e)
77        exit_code = 1
78
79    if parse_common_args().timing_info:
80        print_timing_info()
81
82    sys.exit(exit_code)
83
84
85if __name__ == "__main__":
86    import doctest
87
88    (failures, num_tests) = doctest.testmod(optionflags=doctest.ELLIPSIS)
89    sys.exit(1 if failures > 0 else 0)
90