155 lines · python
1# DExTer : Debugging Experience Tester2# ~~~~~~ ~ ~~ ~ ~~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-exception7"""Extended Argument Parser. Extends the argparse module with some extra8functionality, to hopefully aid user-friendliness.9"""10 11import argparse12import difflib13import unittest14 15from dex.utils import PrettyOutput16from dex.utils.Exceptions import Error17 18# re-export all of argparse19for 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_mean31 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-builtin48 def _print_message(self, message, file=None):49 if message:50 if file and file.name == "<stdout>":51 file = PrettyOutput.stdout52 else:53 file = PrettyOutput.stderr54 55 self.context.o.auto(message, file)56 57 # pylint: enable=redefined-builtin58 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 @property66 def _valid_visible_options(self):67 """A list of all non-suppressed command line flags."""68 return [69 item70 for sublist in vars(self)["_actions"]71 for item in sublist.option_strings72 if sublist.help != argparse.SUPPRESS73 ]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 args92 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 default102 and isinstance(default, (str, int, float))103 and default != argparse.SUPPRESS104 ):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 = context118 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