|
| 1 | +"""Some useful argument types. |
| 2 | +
|
| 3 | +Note that these types can be used with other argparse-compatible libraries, including |
| 4 | +"argparse" itself. |
| 5 | +
|
| 6 | +The 'type' parameter to ArgumentParser.add_argument() must be a callable object, |
| 7 | +typically a function. That function is called to convert the string to the Python type available |
| 8 | +in the 'namespace' passed to your "do_xyz" command function. Thus, "type=int" works because |
| 9 | +int("53") returns the integer value 53. If that callable object / function raises an exception |
| 10 | +due to invalid input, the name ("repr") of the object/function will be printed in the error message |
| 11 | +to the user. Using lambda, functools.partial, or the like will generate a callable object with a |
| 12 | +rather opaque repr so it can be useful to have a one-line function rather than relying on a lambda, |
| 13 | +even for a short expression. |
| 14 | +
|
| 15 | +For "types" that have some context/state, using a class with a __call__ method, and overriding |
| 16 | +the __repr__ method, allows you to produce an error message that provides that information |
| 17 | +to the user. |
| 18 | +""" |
| 19 | + |
| 20 | +from collections.abc import Iterable |
| 21 | + |
| 22 | +import cmd2 |
| 23 | + |
| 24 | +_int_suffixes = { |
| 25 | + # SI number suffixes (unit prefixes): |
| 26 | + "K": 1_000, |
| 27 | + "M": 1_000_000, |
| 28 | + "G": 1_000_000_000, |
| 29 | + "T": 1_000_000_000_000, |
| 30 | + "P": 1_000_000_000_000_000, |
| 31 | + # IEC number suffixes (unit prefixes): |
| 32 | + "Ki": 1024, |
| 33 | + "Mi": 1024 * 1024, |
| 34 | + "Gi": 1024 * 1024 * 1024, |
| 35 | + "Ti": 1024 * 1024 * 1024 * 1024, |
| 36 | + "Pi": 1024 * 1024 * 1024 * 1024 * 1024, |
| 37 | +} |
| 38 | + |
| 39 | + |
| 40 | +def integer(value_str: str) -> int: |
| 41 | + """Will accept any base, and optional suffix like '64K'.""" |
| 42 | + multiplier = 1 |
| 43 | + # If there is a matching suffix, use its multiplier: |
| 44 | + for suffix, suffix_multiplier in _int_suffixes.items(): |
| 45 | + if value_str.endswith(suffix): |
| 46 | + value_str = value_str.removesuffix(suffix) |
| 47 | + multiplier = suffix_multiplier |
| 48 | + break |
| 49 | + |
| 50 | + return int(value_str, 0) * multiplier |
| 51 | + |
| 52 | + |
| 53 | +def hexadecimal(value_str: str) -> int: |
| 54 | + """Parse hexidecimal integer, with optional '0x' prefix.""" |
| 55 | + return int(value_str, base=16) |
| 56 | + |
| 57 | + |
| 58 | +class Range: |
| 59 | + """Useful as type for large ranges, when 'choices=range(maxval)' would be excessively large.""" |
| 60 | + |
| 61 | + def __init__(self, firstval: int, secondval: int | None = None) -> None: |
| 62 | + """Construct a Range, with same syntax as 'range'. |
| 63 | +
|
| 64 | + :param firstval: either the top end of range (if 'secondval' is missing), or the bottom end |
| 65 | + :param secondval: top end of range (one higher than maximum value) |
| 66 | + """ |
| 67 | + if secondval is None: |
| 68 | + self.bottom = 0 |
| 69 | + self.top = firstval |
| 70 | + else: |
| 71 | + self.bottom = firstval |
| 72 | + self.top = secondval |
| 73 | + |
| 74 | + self.range_str = f"[{self.bottom}..{self.top - 1}]" |
| 75 | + |
| 76 | + def __repr__(self) -> str: |
| 77 | + """Will be printed as the 'argument type' to user on syntax or range error.""" |
| 78 | + return f"Range{self.range_str}" |
| 79 | + |
| 80 | + def __call__(self, arg: str) -> int: |
| 81 | + """Parse the string argument and checks validity.""" |
| 82 | + val = integer(arg) |
| 83 | + if self.bottom <= val < self.top: |
| 84 | + return val |
| 85 | + raise ValueError(f"Value '{val}' not within {self.range_str}") |
| 86 | + |
| 87 | + |
| 88 | +class IntSet: |
| 89 | + """Set of integers from a specified range. |
| 90 | +
|
| 91 | + e.g. '5', '1-3,8', 'all' |
| 92 | + """ |
| 93 | + |
| 94 | + def __init__(self, firstval: int, secondval: int | None = None) -> None: |
| 95 | + """Construct an IntSet, with same syntax as 'range'. |
| 96 | +
|
| 97 | + :param firstval: either the top end of range (if 'secondval' is missing), or the bottom end |
| 98 | + :param secondval: top end of range (one higher than maximum value) |
| 99 | + """ |
| 100 | + if secondval is None: |
| 101 | + self.bottom = 0 |
| 102 | + self.top = firstval |
| 103 | + else: |
| 104 | + self.bottom = firstval |
| 105 | + self.top = secondval |
| 106 | + |
| 107 | + self.range_str = f"[{self.bottom}..{self.top - 1}]" |
| 108 | + |
| 109 | + def __repr__(self) -> str: |
| 110 | + """Will be printed as the 'argument type' to user on syntax or range error.""" |
| 111 | + return f"IntSet{self.range_str}" |
| 112 | + |
| 113 | + def __call__(self, arg: str) -> Iterable[int]: |
| 114 | + """Parse a string into an iterable returning ints.""" |
| 115 | + if arg == 'all': |
| 116 | + return range(self.bottom, self.top) |
| 117 | + |
| 118 | + out = [] |
| 119 | + for piece in arg.split(','): |
| 120 | + if '-' in piece: |
| 121 | + a, b = [int(x) for x in piece.split('-', 2)] |
| 122 | + if a < self.bottom: |
| 123 | + raise ValueError(f"Value '{a}' not within {self.range_str}") |
| 124 | + if b >= self.top: |
| 125 | + raise ValueError(f"Value '{b}' not within {self.range_str}") |
| 126 | + out += list(range(a, b + 1)) |
| 127 | + else: |
| 128 | + val = int(piece) |
| 129 | + if not self.bottom <= val < self.top: |
| 130 | + raise ValueError(f"Value '{val}' not within {self.range_str}") |
| 131 | + out += [val] |
| 132 | + return out |
| 133 | + |
| 134 | + |
| 135 | +if __name__ == '__main__': |
| 136 | + import argparse |
| 137 | + import sys |
| 138 | + |
| 139 | + class CustomTypesExample(cmd2.Cmd): |
| 140 | + example_parser = cmd2.Cmd2ArgumentParser() |
| 141 | + example_parser.add_argument( |
| 142 | + '--value', '-v', type=integer, help='Integer value, with optional K/M/G/Ki/Mi/Gi/... suffix' |
| 143 | + ) |
| 144 | + example_parser.add_argument('--memory-address', '-m', type=hexadecimal, help='Memory address in hex') |
| 145 | + example_parser.add_argument('--year', type=Range(1900, 2000), help='Year between 1900-1999') |
| 146 | + example_parser.add_argument( |
| 147 | + '--index', dest='index_list', type=IntSet(100), help='One or more indexes 0-99. e.g. "1,3,5", "10,30-50", "all"' |
| 148 | + ) |
| 149 | + |
| 150 | + @cmd2.with_argparser(example_parser) |
| 151 | + def do_example(self, args: argparse.Namespace) -> None: |
| 152 | + """The example command.""" |
| 153 | + if args.value is not None: |
| 154 | + self.poutput(f"Value: {args.value}") |
| 155 | + if args.memory_address is not None: |
| 156 | + # print the value as hex, with leading "0x" + 16 hex digits + three '_' group separators: |
| 157 | + self.poutput(f"Address: {args.memory_address:#021_x}") |
| 158 | + if args.year is not None: |
| 159 | + self.poutput(f"Year: {args.year}") |
| 160 | + if args.index_list is not None: |
| 161 | + for index in args.index_list: |
| 162 | + self.poutput(f"Process index {index}") |
| 163 | + |
| 164 | + app = CustomTypesExample() |
| 165 | + sys.exit(app.cmdloop()) |
0 commit comments