File: custom_types.py

package info (click to toggle)
cmd2 3.2.0%2Bds-1
  • links: PTS, VCS
  • area: main
  • in suites: experimental
  • size: 2,664 kB
  • sloc: python: 17,488; makefile: 114; sh: 39; javascript: 7
file content (165 lines) | stat: -rw-r--r-- 6,256 bytes parent folder | download
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
"""Some useful argument types.

Note that these types can be used with other argparse-compatible libraries, including
"argparse" itself.

The 'type' parameter to ArgumentParser.add_argument() must be a callable object,
typically a function. That function is called to convert the string to the Python type available
in the 'namespace' passed to your "do_xyz" command function. Thus, "type=int" works because
int("53") returns the integer value 53. If that callable object / function raises an exception
due to invalid input, the name ("repr") of the object/function will be printed in the error message
to the user. Using lambda, functools.partial, or the like will generate a callable object with a
rather opaque repr so it can be useful to have a one-line function rather than relying on a lambda,
even for a short expression.

For "types" that have some context/state, using a class with a __call__ method, and overriding
the __repr__ method, allows you to produce an error message that provides that information
to the user.
"""

from collections.abc import Iterable

import cmd2

_int_suffixes = {
    # SI number suffixes (unit prefixes):
    "K": 1_000,
    "M": 1_000_000,
    "G": 1_000_000_000,
    "T": 1_000_000_000_000,
    "P": 1_000_000_000_000_000,
    # IEC number suffixes (unit prefixes):
    "Ki": 1024,
    "Mi": 1024 * 1024,
    "Gi": 1024 * 1024 * 1024,
    "Ti": 1024 * 1024 * 1024 * 1024,
    "Pi": 1024 * 1024 * 1024 * 1024 * 1024,
}


def integer(value_str: str) -> int:
    """Will accept any base, and optional suffix like '64K'."""
    multiplier = 1
    # If there is a matching suffix, use its multiplier:
    for suffix, suffix_multiplier in _int_suffixes.items():
        if value_str.endswith(suffix):
            value_str = value_str.removesuffix(suffix)
            multiplier = suffix_multiplier
            break

    return int(value_str, 0) * multiplier


def hexadecimal(value_str: str) -> int:
    """Parse hexidecimal integer, with optional '0x' prefix."""
    return int(value_str, base=16)


class Range:
    """Useful as type for large ranges, when 'choices=range(maxval)' would be excessively large."""

    def __init__(self, firstval: int, secondval: int | None = None) -> None:
        """Construct a Range, with same syntax as 'range'.

        :param firstval: either the top end of range (if 'secondval' is missing), or the bottom end
        :param secondval: top end of range (one higher than maximum value)
        """
        if secondval is None:
            self.bottom = 0
            self.top = firstval
        else:
            self.bottom = firstval
            self.top = secondval

        self.range_str = f"[{self.bottom}..{self.top - 1}]"

    def __repr__(self) -> str:
        """Will be printed as the 'argument type' to user on syntax or range error."""
        return f"Range{self.range_str}"

    def __call__(self, arg: str) -> int:
        """Parse the string argument and checks validity."""
        val = integer(arg)
        if self.bottom <= val < self.top:
            return val
        raise ValueError(f"Value '{val}' not within {self.range_str}")


class IntSet:
    """Set of integers from a specified range.

    e.g. '5', '1-3,8', 'all'
    """

    def __init__(self, firstval: int, secondval: int | None = None) -> None:
        """Construct an IntSet, with same syntax as 'range'.

        :param firstval: either the top end of range (if 'secondval' is missing), or the bottom end
        :param secondval: top end of range (one higher than maximum value)
        """
        if secondval is None:
            self.bottom = 0
            self.top = firstval
        else:
            self.bottom = firstval
            self.top = secondval

        self.range_str = f"[{self.bottom}..{self.top - 1}]"

    def __repr__(self) -> str:
        """Will be printed as the 'argument type' to user on syntax or range error."""
        return f"IntSet{self.range_str}"

    def __call__(self, arg: str) -> Iterable[int]:
        """Parse a string into an iterable returning ints."""
        if arg == 'all':
            return range(self.bottom, self.top)

        out = []
        for piece in arg.split(','):
            if '-' in piece:
                a, b = [int(x) for x in piece.split('-', 2)]
                if a < self.bottom:
                    raise ValueError(f"Value '{a}' not within {self.range_str}")
                if b >= self.top:
                    raise ValueError(f"Value '{b}' not within {self.range_str}")
                out += list(range(a, b + 1))
            else:
                val = int(piece)
                if not self.bottom <= val < self.top:
                    raise ValueError(f"Value '{val}' not within {self.range_str}")
                out += [val]
        return out


if __name__ == '__main__':
    import argparse
    import sys

    class CustomTypesExample(cmd2.Cmd):
        example_parser = cmd2.Cmd2ArgumentParser()
        example_parser.add_argument(
            '--value', '-v', type=integer, help='Integer value, with optional K/M/G/Ki/Mi/Gi/... suffix'
        )
        example_parser.add_argument('--memory-address', '-m', type=hexadecimal, help='Memory address in hex')
        example_parser.add_argument('--year', type=Range(1900, 2000), help='Year between 1900-1999')
        example_parser.add_argument(
            '--index', dest='index_list', type=IntSet(100), help='One or more indexes 0-99. e.g. "1,3,5", "10,30-50", "all"'
        )

        @cmd2.with_argparser(example_parser)
        def do_example(self, args: argparse.Namespace) -> None:
            """The example command."""
            if args.value is not None:
                self.poutput(f"Value: {args.value}")
            if args.memory_address is not None:
                # print the value as hex, with leading "0x" + 16 hex digits + three '_' group separators:
                self.poutput(f"Address: {args.memory_address:#021_x}")
            if args.year is not None:
                self.poutput(f"Year: {args.year}")
            if args.index_list is not None:
                for index in args.index_list:
                    self.poutput(f"Process index {index}")

    app = CustomTypesExample()
    sys.exit(app.cmdloop())