~brenns10/subc

subc/subc/__init__.py -rw-r--r-- 9.7 KiB
12881154Stephen Brennan Release v0.7.1 2 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
#!/usr/bin/env python3
"""
A simple sub-command library for writing rich CLIs
"""
import argparse
import collections
import typing as t
from abc import ABC
from abc import abstractproperty
from abc import abstractmethod


def _first_different(s1: str, s2: str) -> int:
    """
    Return index of the first different character in s1 or s2. If the strings
    are the same, raises a ValueError.
    """
    for i, (c1, c2) in enumerate(zip(s1, s2)):
        if c1 != c2:
            return i
    if len(s1) == len(s2):
        raise ValueError(f"Duplicate string {s1!r} is not allowed")
    return i + 1


def _unique_prefixes(strings: t.Iterable[str]) -> t.Dict[str, t.List[str]]:
    """
    Helper to find a list of unique prefixes for each string in strings.

    Return a dict mapping each string to a list of prefixes which are unique
    among all other strings within the list. Here is an example:

        >>> _unique_prefixes(["commit", "count", "apply", "app", "shape"])
        {'app': [],
         'apply': ['appl'],
         'commit': ['com', 'comm', 'commi'],
         'count': ['cou', 'coun'],
         'launch': ['la', 'lau', 'laun', 'launc'],
         'list': ['li', 'lis'],
         'shape': ['s', 'sh', 'sha', 'shap']}
    """
    strings = sorted(strings)
    diffs = [0] * len(strings)
    for i, (s1, s2) in enumerate(zip(strings, strings[1:])):
        common = _first_different(s1, s2)
        diffs[i] = max(diffs[i], common)
        diffs[i + 1] = max(diffs[i + 1], common)
    return {
        s: [s[:i] for i in range(x + 1, len(s))]
        for (s, x) in zip(strings, diffs)
    }


class _SneakyDict(collections.UserDict):
    """
    A dictionary which can have "hidden" keys that only show up if you know
    about them. The keys are just aliases to other keys. They show up with
    "getitem" and "contains" operations, but not in list / len operations.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._aliases = {}

    def __getitem__(self, key):
        key = self._aliases.get(key, key)
        return super().__getitem__(key)

    def __contains__(self, key):
        key = self._aliases.get(key, key)
        return super().__contains__(key)

    def add_aliases(self, alias_map: t.Dict[str, t.List[str]]):
        alias_to_name = {a: n for n, l in alias_map.items() for a in l}
        self._aliases.update(alias_to_name)


def _wrap_subparser_aliases(
        option: argparse._SubParsersAction,
        alias_map: t.Dict[str, t.List[str]]
) -> None:
    """
    Unfortunately, this mucks around with an internal implementation of
    argparse. However, the API seems pretty stable, and I hope to catch any
    compatibility issues with testing on each new version.

    The "choices" and "_name_parser_map" fields are used to determine which
    subcommands are allowed, and also to list out all of the subcommands for the
    help output (or even to generate completions with something like
    argcomplete).

    For the purposes of lookup (or membership testing), we want the aliases to
    be reflected in these variables. But for the purposes of listing, the
    aliases should be hidden. Thus, use a the _SneakyDict from above to hide the
    aliases.
    """
    new_choices = _SneakyDict(option.choices)
    new_choices.add_aliases(alias_map)
    option.choices = new_choices  # type: ignore
    option._name_parser_map = option.choices


class Command(ABC):
    """
    A simple class for implementing sub-commands in your command line
    application. Create a subclass for your app as follows:

        class MyCmd(subc.Command):
            pass

    Then, each command in your app can subclass this, implementing the three
    required fields:

        class HelloWorld(MyCmd):
            name = 'hello-world'
            description = 'say hello'
            def run(self):
                print('hello world')

    Finally, use your app-level subclass for creating an argument parser:

        def main():
            parser = argparse.ArgumentParser(description='a cool tool')
            MyCmd.add_commands(parser)
            args = parser.parse_args()
            args.func(args)
    """

    @abstractproperty
    def name(self) -> str:
        """A field or property which is used for the command name argument"""

    @abstractproperty
    def description(self) -> str:
        """A field or property which is used as the help/description"""

    def add_args(self, parser: argparse.ArgumentParser):
        pass  # default is no arguments

    @abstractmethod
    def run(self) -> t.Any:
        """Function which is called for this command."""

    def base_run(self, args: argparse.Namespace):
        self.args = args
        return self.run()

    @classmethod
    def add_commands(
        cls,
        parser: argparse.ArgumentParser,
        default: t.Optional[str] = None,
        shortest_prefix: bool = False,
        cmd_aliases: t.Optional[t.Mapping[str, str]] = None,
    ) -> argparse.ArgumentParser:
        """
        Add all subcommands which are descendents of this class to parser.

        This call is required in order to setup an argument parser before
        parsing args and executing sub-command. Each sub-command must be a
        sub-class (or a further descendent) of this class. Only leaf subclasses
        are considered commands -- internal "nodes" in the hierarchy are skipped
        as they are assumed to be helpers.

        A default command to run may be set with 'default'. When the argument
        parser is called without a sub-command, this command will automatically
        execute (rather than simply raising an Exception).

        Shortest prefix sub-command matching allows the user to select a
        sub-command by using any string which is a prefix of exactly one
        command, e.g. "git cl" rather than "git clone". This is useful whenever
        there is a small, unchanging set of sub-commands, as a user can develop
        muscle memory for prefixes. However, if the set of sub-commands changes
        over time, then users may develop muscle-memory for a prefix which
        becomes ambiguous with a new command. Thus, it may be preferable to
        allow users to specify their own alias list. This function supports
        both methods of aliasing, or no aliasing.

        :param parser: Argument parser which is already created for this app
        :param default: Name of the command which should be executed if none is
          selected
        :param shortest_prefix: Enable shortest prefix command matching. If
          enabled, this takes priority over cmd_aliases.
        :param cmd_aliases: User-provided alias list in the form
          {"alias": "true name"}. This is only used if shortest_prefix is
          False.
        :returns: the modified parser (this can be ignored)
        """
        default_set = False
        subparsers = parser.add_subparsers(title='sub-command')
        subclasses = collections.deque(cls.__subclasses__())
        to_add = []
        while subclasses:
            subcls = subclasses.popleft()
            this_node_subclasses = subcls.__subclasses__()
            if this_node_subclasses:
                # Assume that any class with children is not executable. Add
                # its children to the queue (BFS) but do not instantiate it.
                subclasses.extend(this_node_subclasses)
            else:
                to_add.append(subcls())

        if shortest_prefix:
            aliases = _unique_prefixes(c.name for c in to_add)
        elif cmd_aliases:
            aliases = collections.defaultdict(list)
            for name, target in cmd_aliases.items():
                aliases[target].append(name)
        else:
            aliases = collections.defaultdict(list)
        for cmd in to_add:
            cmd_parser = subparsers.add_parser(
                cmd.name, description=cmd.description,
                help=getattr(cmd, "help", cmd.description),
            )
            cmd.add_args(cmd_parser)
            cmd_parser.set_defaults(func=cmd.base_run)
            if cmd.name == default:
                parser.set_defaults(func=cmd.base_run)
                default_set = True
        _wrap_subparser_aliases(subparsers, aliases)

        if not default_set:
            def default_func(*args, **kwargs):
                raise Exception('you must select a sub-command')
            parser.set_defaults(func=default_func)
        return parser

    @classmethod
    def main(
            cls,
            description: str,
            default: t.Optional[str] = None,
            args: t.Optional[t.List[str]] = None,
            shortest_prefix: bool = False,
    ) -> t.Any:
        """
        Parse arguments and run the selected sub-command.

        This helper function is expected to be the main, most useful API for
        subc, although you could directly call the add_commands() method.
        Creates an argument parser, adds every discovered sub-command, parses
        the arguments, and executes the selected sub-command, returning its
        return value.

        Custom arguments (rather than sys.argv) can be specified using "args".
        Details on the arguments "default" and "shortest_prefix" can be found
        in the docstring for add_commands().

        :param description: Description of the application (for help output)
        :param default: Default command name
        :param args: If specified, a list of args to use in place of sys.argv
        :param shortest_prefix: whether to enable prefix matching
        :returns: Return value of the selected command's run() method
        """
        parser = argparse.ArgumentParser(description=description)
        cls.add_commands(
            parser, default=default, shortest_prefix=shortest_prefix,
        )
        ns = parser.parse_args(args=args)
        return ns.func(ns)