xref: /llvm-project/cross-project-tests/debuginfo-tests/dexter/dex/utils/ExtArgParse.py (revision f98ee40f4b5d7474fc67e82824bf6abbaedb7b1c)
1# DExTer : Debugging Experience Tester
2# ~~~~~~   ~         ~~         ~   ~~
3#
4# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5# See https://llvm.org/LICENSE.txt for license information.
6# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7"""Extended Argument Parser. Extends the argparse module with some extra
8functionality, to hopefully aid user-friendliness.
9"""
10
11import argparse
12import difflib
13import unittest
14
15from dex.utils import PrettyOutput
16from dex.utils.Exceptions import Error
17
18# re-export all of argparse
19for argitem in argparse.__all__:
20    vars()[argitem] = getattr(argparse, argitem)
21
22
23def _did_you_mean(val, possibles):
24    close_matches = difflib.get_close_matches(val, possibles)
25    did_you_mean = ""
26    if close_matches:
27        did_you_mean = "did you mean {}?".format(
28            " or ".join("<y>'{}'</>".format(c) for c in close_matches[:2])
29        )
30    return did_you_mean
31
32
33def _colorize(message):
34    lines = message.splitlines()
35    for i, line in enumerate(lines):
36        lines[i] = lines[i].replace("usage:", "<g>usage:</>")
37        if line.endswith(":"):
38            lines[i] = "<g>{}</>".format(line)
39    return "\n".join(lines)
40
41
42class ExtArgumentParser(argparse.ArgumentParser):
43    def error(self, message):
44        """Use the Dexception Error mechanism (including auto-colored output)."""
45        raise Error("{}\n\n{}".format(message, self.format_usage()))
46
47    # pylint: disable=redefined-builtin
48    def _print_message(self, message, file=None):
49        if message:
50            if file and file.name == "<stdout>":
51                file = PrettyOutput.stdout
52            else:
53                file = PrettyOutput.stderr
54
55            self.context.o.auto(message, file)
56
57    # pylint: enable=redefined-builtin
58
59    def format_usage(self):
60        return _colorize(super(ExtArgumentParser, self).format_usage())
61
62    def format_help(self):
63        return _colorize(super(ExtArgumentParser, self).format_help() + "\n\n")
64
65    @property
66    def _valid_visible_options(self):
67        """A list of all non-suppressed command line flags."""
68        return [
69            item
70            for sublist in vars(self)["_actions"]
71            for item in sublist.option_strings
72            if sublist.help != argparse.SUPPRESS
73        ]
74
75    def parse_args(self, args=None, namespace=None):
76        """Add 'did you mean' output to errors."""
77        args, argv = self.parse_known_args(args, namespace)
78        if argv:
79            errors = []
80            for arg in argv:
81                if arg in self._valid_visible_options:
82                    error = "unexpected argument: <y>'{}'</>".format(arg)
83                else:
84                    error = "unrecognized argument: <y>'{}'</>".format(arg)
85                    dym = _did_you_mean(arg, self._valid_visible_options)
86                    if dym:
87                        error += "  ({})".format(dym)
88                errors.append(error)
89            self.error("\n       ".join(errors))
90
91        return args
92
93    def add_argument(self, *args, **kwargs):
94        """Automatically add the default value to help text."""
95        if "default" in kwargs:
96            default = kwargs["default"]
97            if default is None:
98                default = kwargs.pop("display_default", None)
99
100            if (
101                default
102                and isinstance(default, (str, int, float))
103                and default != argparse.SUPPRESS
104            ):
105                assert (
106                    "choices" not in kwargs or default in kwargs["choices"]
107                ), "default value '{}' is not one of allowed choices: {}".format(
108                    default, kwargs["choices"]
109                )
110                if "help" in kwargs and kwargs["help"] != argparse.SUPPRESS:
111                    assert isinstance(kwargs["help"], str), type(kwargs["help"])
112                    kwargs["help"] = "{} (default:{})".format(kwargs["help"], default)
113
114        super(ExtArgumentParser, self).add_argument(*args, **kwargs)
115
116    def __init__(self, context, *args, **kwargs):
117        self.context = context
118        super(ExtArgumentParser, self).__init__(*args, **kwargs)
119
120
121class TestExtArgumentParser(unittest.TestCase):
122    def test_did_you_mean(self):
123        parser = ExtArgumentParser(None)
124        parser.add_argument("--foo")
125        parser.add_argument("--qoo", help=argparse.SUPPRESS)
126        parser.add_argument("jam", nargs="?")
127
128        parser.parse_args(["--foo", "0"])
129
130        expected = (
131            r"^unrecognized argument\: <y>'\-\-doo'</>\s+"
132            r"\(did you mean <y>'\-\-foo'</>\?\)\n"
133            r"\s*<g>usage:</>"
134        )
135        with self.assertRaisesRegex(Error, expected):
136            parser.parse_args(["--doo"])
137
138        parser.add_argument("--noo")
139
140        expected = (
141            r"^unrecognized argument\: <y>'\-\-doo'</>\s+"
142            r"\(did you mean <y>'\-\-noo'</> or <y>'\-\-foo'</>\?\)\n"
143            r"\s*<g>usage:</>"
144        )
145        with self.assertRaisesRegex(Error, expected):
146            parser.parse_args(["--doo"])
147
148        expected = r"^unrecognized argument\: <y>'\-\-bar'</>\n" r"\s*<g>usage:</>"
149        with self.assertRaisesRegex(Error, expected):
150            parser.parse_args(["--bar"])
151
152        expected = r"^unexpected argument\: <y>'\-\-foo'</>\n" r"\s*<g>usage:</>"
153        with self.assertRaisesRegex(Error, expected):
154            parser.parse_args(["--", "x", "--foo"])
155