Skip to content

Commit dfc7fbf

Browse files
committed
Add 'argtypes.py', with useful argument types for users
- integer: like 'int', but accepts arbitrary bases (i.e. 0x10), and optional suffix like "10K" or "64Mi" - hexadecimal: base 16 integer, with nicer error text than using a lambda - Range: integer from specified range, when that may be too large for 'choices' - IntSet: set of integers from specified range
1 parent 21657a6 commit dfc7fbf

File tree

2 files changed

+242
-0
lines changed

2 files changed

+242
-0
lines changed

cmd2/argtypes.py

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
"""Useful argument types."""
2+
3+
from collections.abc import Iterable
4+
5+
_int_suffixes = {
6+
# SI number suffixes (unit prefixes):
7+
"K": 1_000,
8+
"M": 1_000_000,
9+
"G": 1_000_000_000,
10+
"T": 1_000_000_000_000,
11+
"P": 1_000_000_000_000_000,
12+
# IEC number suffixes (unit prefixes):
13+
"Ki": 1024,
14+
"Mi": 1024 * 1024,
15+
"Gi": 1024 * 1024 * 1024,
16+
"Ti": 1024 * 1024 * 1024 * 1024,
17+
"Pi": 1024 * 1024 * 1024 * 1024 * 1024,
18+
}
19+
20+
21+
def integer(value_str: str) -> int:
22+
"""Will accept any base, and optional suffix like '64K'."""
23+
multiplier = 1
24+
# If there is a matching suffix, use its multiplier:
25+
for suffix, suffix_multiplier in _int_suffixes.items():
26+
if value_str.endswith(suffix):
27+
value_str = value_str.removesuffix(suffix)
28+
multiplier = suffix_multiplier
29+
break
30+
31+
return int(value_str, 0) * multiplier
32+
33+
34+
def hexadecimal(value_str: str) -> int:
35+
"""Parse hexidecimal integer, with optional '0x' prefix."""
36+
return int(value_str, base=16)
37+
38+
39+
class Range:
40+
"""Useful as type for large ranges, when 'choices=range(maxval)' would be excessively large."""
41+
42+
def __init__(self, firstval: int, secondval: int | None = None) -> None:
43+
"""Construct a Range, with same syntax as 'range'.
44+
45+
:param firstval: either the top end of range (if 'secondval' is missing), or the bottom end
46+
:param secondval: top end of range (one higher than maximum value)
47+
"""
48+
if secondval is None:
49+
self.bottom = 0
50+
self.top = firstval
51+
else:
52+
self.bottom = firstval
53+
self.top = secondval
54+
55+
self.range_str = f"[{self.bottom}..{self.top - 1}]"
56+
57+
def __repr__(self) -> str:
58+
"""Will be printed as the 'argument type' to user on syntax or range error."""
59+
return f"Range{self.range_str}"
60+
61+
def __call__(self, arg: str) -> int:
62+
"""Parse the string argument and checks validity."""
63+
val = integer(arg)
64+
if self.bottom <= val < self.top:
65+
return val
66+
raise ValueError(f"Value '{val}' not within {self.range_str}")
67+
68+
69+
class IntSet:
70+
"""Set of integers from a specified range.
71+
72+
e.g. '5', '1-3,8', 'all'
73+
"""
74+
75+
def __init__(self, firstval: int, secondval: int | None = None) -> None:
76+
"""Construct an IntSet, with same syntax as 'range'.
77+
78+
:param firstval: either the top end of range (if 'secondval' is missing), or the bottom end
79+
:param secondval: top end of range (one higher than maximum value)
80+
"""
81+
if secondval is None:
82+
self.bottom = 0
83+
self.top = firstval
84+
else:
85+
self.bottom = firstval
86+
self.top = secondval
87+
88+
self.range_str = f"[{self.bottom}..{self.top - 1}]"
89+
90+
def __repr__(self) -> str:
91+
"""Will be printed as the 'argument type' to user on syntax or range error."""
92+
return f"IntSet{self.range_str}"
93+
94+
def __call__(self, arg: str) -> Iterable:
95+
"""Parse a string into an iterable returning ints."""
96+
if arg == 'all':
97+
return range(self.bottom, self.top)
98+
99+
out = []
100+
for piece in arg.split(','):
101+
if '-' in piece:
102+
a, b = [int(x) for x in piece.split('-', 2)]
103+
if a < self.bottom:
104+
raise ValueError(f"Value '{a}' not within {self.range_str}")
105+
if b >= self.top:
106+
raise ValueError(f"Value '{b}' not within {self.range_str}")
107+
out += list(range(a, b + 1))
108+
else:
109+
val = int(piece)
110+
if not self.bottom <= val < self.top:
111+
raise ValueError(f"Value '{val}' not within {self.range_str}")
112+
out += [val]
113+
return out

tests/test_argtypes.py

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
"""Tests of additional argument types"""
2+
3+
import argparse
4+
5+
import pytest
6+
7+
import cmd2
8+
from cmd2.argtypes import IntSet, Range, hexadecimal, integer
9+
10+
from .conftest import (
11+
run_cmd,
12+
)
13+
14+
15+
class CustomTypesApp(cmd2.Cmd):
16+
def __init__(self) -> None:
17+
cmd2.Cmd.__init__(self)
18+
19+
def namespace_provider(self) -> argparse.Namespace:
20+
ns = argparse.Namespace()
21+
ns.custom_stuff = "custom"
22+
return ns
23+
24+
@staticmethod
25+
def _test_parser_builder() -> cmd2.Cmd2ArgumentParser:
26+
test_parser = cmd2.Cmd2ArgumentParser()
27+
test_parser.add_argument('--int', dest='intarg', type=integer, help='IntegerHelp')
28+
test_parser.add_argument('--hex', dest='hexarg', type=hexadecimal, help='HexHelp')
29+
test_parser.add_argument('--range', dest='rangearg', type=Range(10), help='RangeHelp')
30+
test_parser.add_argument('--set', dest='setarg', type=IntSet(5), help='SetHelp')
31+
test_parser.add_argument('--highset', dest='setarg', type=IntSet(5, 10), help='SetHelp')
32+
return test_parser
33+
34+
@cmd2.with_argparser(_test_parser_builder)
35+
def do_test(self, args, *, keyword_arg: str | None = None) -> None:
36+
"""Test custom types
37+
:param args: argparse namespace
38+
:param keyword_arg: Optional keyword arguments
39+
"""
40+
if args.intarg is not None:
41+
self.stdout.write(f"Integer {args.intarg}")
42+
if args.hexarg is not None:
43+
print(f"Hex {args.hexarg}")
44+
if args.rangearg is not None:
45+
print(f"Range {args.rangearg}")
46+
if args.setarg is not None:
47+
print(f"Set {list(args.setarg)}")
48+
49+
50+
@pytest.fixture
51+
def custom_types_app():
52+
return CustomTypesApp()
53+
54+
55+
def test_int_basic(custom_types_app) -> None:
56+
out, _err = run_cmd(custom_types_app, 'test --int 5')
57+
assert out == ['Integer 5']
58+
59+
60+
def test_int_hex(custom_types_app) -> None:
61+
out, _err = run_cmd(custom_types_app, 'test --int 0xf')
62+
assert out == ['Integer 15']
63+
64+
65+
def test_int_bin(custom_types_app) -> None:
66+
out, _err = run_cmd(custom_types_app, 'test --int 0b1000_0000')
67+
assert out == ['Integer 128']
68+
69+
70+
def test_int_si_suffix(custom_types_app) -> None:
71+
out, _err = run_cmd(custom_types_app, 'test --int 5M')
72+
assert out == ['Integer 5000000']
73+
74+
75+
def test_int_iec_suffix(custom_types_app) -> None:
76+
out, _err = run_cmd(custom_types_app, 'test --int 10Ki')
77+
assert out == ['Integer 10240']
78+
79+
80+
def test_int_failure(custom_types_app) -> None:
81+
_out, err = run_cmd(custom_types_app, 'test --int 5bob')
82+
assert err[-1:] == ["Error: argument --int: invalid integer value: '5bob'"]
83+
84+
85+
def test_hex_basic(custom_types_app) -> None:
86+
out, _err = run_cmd(custom_types_app, 'test --hex 10')
87+
assert out == ['Hex 16']
88+
89+
90+
def test_hex_prefixed(custom_types_app) -> None:
91+
out, _err = run_cmd(custom_types_app, 'test --hex 0x10')
92+
assert out == ['Hex 16']
93+
94+
95+
def test_hex_failure(custom_types_app) -> None:
96+
_out, err = run_cmd(custom_types_app, 'test --hex 5bob')
97+
assert err[-1:] == ["Error: argument --hex: invalid hexadecimal value: '5bob'"]
98+
99+
100+
def test_range(custom_types_app) -> None:
101+
out, _err = run_cmd(custom_types_app, 'test --range 4')
102+
assert out == ['Range 4']
103+
104+
105+
def test_range_outside(custom_types_app) -> None:
106+
_out, err = run_cmd(custom_types_app, 'test --range 10')
107+
assert err[-1:] == ["Error: argument --range: invalid Range[0..9] value: '10'"]
108+
109+
110+
def test_intset_single(custom_types_app) -> None:
111+
out, _err = run_cmd(custom_types_app, 'test --set 4')
112+
assert out == ['Set [4]']
113+
114+
115+
def test_intset_single_out_of_range(custom_types_app) -> None:
116+
_out, err = run_cmd(custom_types_app, 'test --set 5')
117+
assert err[-1:] == ["Error: argument --set: invalid IntSet[0..4] value: '5'"]
118+
119+
120+
def test_intset_high_multi(custom_types_app) -> None:
121+
out, _err = run_cmd(custom_types_app, 'test --highset 5,7-9')
122+
print(_err)
123+
print(out)
124+
assert out == ['Set [5, 7, 8, 9]']
125+
126+
127+
def test_intset_high_all(custom_types_app) -> None:
128+
out, _err = run_cmd(custom_types_app, 'test --highset all')
129+
assert out == ['Set [5, 6, 7, 8, 9]']

0 commit comments

Comments
 (0)