3 Module which provides extensions for the standard Python argparse module.
4 Copyright 2015 Cray Inc. All Rights Reserved
10 from argparse import _ensure_value, Action, ArgumentTypeError
13 class ExtendAction(Action):
14 """Action to extend a list of argument values
16 This action is similar to the standard AppendAction, but uses the
17 extend() attribute of lists rather than the append() attribute. As
18 such, it also has an additional requirement:
20 - This action must receive an iterable 'values' argument from the
21 parser. There are two ways to make this happen:
23 1. Use type= to produce an iterable, e.g. type=str or type=list
24 2. Use nargs= to cause the parser to produe a list, which it
25 does for any nargs= setting that is not None (default) and
29 def __call__(self, parser, namespace, values, option_string):
30 items = copy.copy(_ensure_value(namespace, self.dest, []))
35 # Assume the TypeError is because values is not iterable
36 raise ArgumentTypeError(
37 "argument type '{:s}' is not iterable".format(
38 type(values).__name__))
40 setattr(namespace, self.dest, items)
43 def str2list(string, sep=',', totype=None, choices=None):
44 """Split a string into a list with conversion and validation
46 Split a string into a list, optionally convert each element to a
47 given type and optionally validate that all resulting values are
48 in a collection of valid values.
51 plural = {False: '', True: 's'}
53 # Values should be string or an iterable container of strings.
54 # Split values on the separator into a list
56 lst = string.split(sep)
57 except AttributeError:
58 raise ArgumentTypeError(
59 "argument type '{:s}' does not have split() attribute".format(
60 type(string).__name__))
62 # Perform type conversion
63 if totype is not None:
65 for i, v in enumerate(lst):
68 except (TypeError, ValueError):
71 msg = "invalid {:s} value{:s}: {!r:s}".format(
72 totype.__name__, plural[len(errs) > 1], errs)
73 raise ArgumentTypeError(msg)
75 # Verify each separate value
76 if choices is not None:
77 errs = filter(lambda x:x not in choices, lst)
79 msg = "invalid choice{:s}: {!r:s} (choose from {!s:s})".format(
80 plural[len(errs) > 1], errs, choices)
81 raise ArgumentTypeError(msg)
86 def tolist(sep=',', totype=None, choices=None):
87 """Returns a parameterized callable for argument parser type conversion
89 This function returns a function which accepts a single argument at
90 call time and which uses the supplied arguments to modify its conversion
93 return lambda x:str2list(x, sep=sep, totype=totype, choices=choices)
96 if __name__ == '__main__':
99 class Test_Action_Base(unittest.TestCase):
100 """Create a base class for testing argparse Action classes"""
103 """Create the ExtendAction object and args Namespace"""
104 self.action = ExtendAction([], dest='dest')
105 self.args = argparse.Namespace()
107 def actionRun(self, values):
108 """Run the Action instance using values"""
109 self.action(None, self.args, values, '')
111 def actionEqual(self, values, expected):
112 """Run the Action and check the expected results"""
113 self.actionRun(values)
114 self.assertEqual(self.args.dest, expected)
116 def actionArgTypeErr(self, values):
117 """Run the Action and verify it raises ArgumentTypeError"""
119 ArgumentTypeError, self.action, None, self.args, values, '')
122 class Test_ExtendAction(Test_Action_Base):
123 """Test the ExtendAction class"""
125 def test_non_iterable(self):
126 """Test ExtendAction with a non-iterable type
129 parser.add_argument('-z', nargs=None, type=int ...)
131 parser.parse_args(['-z', '0'])
133 self.actionArgTypeErr(0)
135 def test_single_value(self):
136 """Test ExtendAction with a single value
139 parser.add_argument('-z', nargs=None ...)
141 parser.parse_args(['-z', 'a'])
143 self.actionEqual('a', ['a'])
145 def test_single_string(self):
146 """Test ExtendAction with a single value
149 parser.add_argument('-z', nargs=None ...)
151 parser.parse_args(['-z', 'abc'])
153 self.actionEqual('abc', ['a', 'b', 'c'])
155 def test_single_value_multiple_calls(self):
156 """Test ExtendAction with a single value and multiple calls
159 parser.add_argument('-z', nargs=None, type=int ...)
161 parser.parse_args(['-z', 'a', '-z', 'b'])
163 self.actionEqual('a', ['a'])
164 self.actionEqual('b', ['a', 'b'])
166 def test_value_list(self):
167 """Test ExtendAction with a value list
170 parser.add_argument('-z', nargs=1 ...)
172 parser.parse_args(['-z', 'abc'])
174 self.actionEqual(['abc'], ['abc'])
176 def test_value_list_multiple_calls(self):
177 """Test ExtendAction with a single value and multiple calls
180 parser.add_argument('-z', nargs=1 ...)
182 parser.parse_args(['-z', 'abc', '-z', 'def'])
184 self.actionRun(['abc'])
185 self.actionEqual(['def'], ['abc', 'def'])
187 def test_value_list_multiple_values(self):
188 """Test ExtendAction with a value list of length > 1
191 parser.add_argument('-z', nargs=2 ...)
193 parser.add_argument('-z', nargs='+' ...)
195 parser.add_argument('-z', nargs='*' ...)
197 parser.parse_args(['-z', 'abc', 'def'])
199 self.actionEqual(['abc', 'def'], ['abc', 'def'])
202 class Test_tolist_str2list(unittest.TestCase):
203 """Test the str2list and tolist conversion functions"""
206 """Verify default and non-default separators work"""
208 self.assertEqual(f('a,b,c'), ['a','b','c'])
210 self.assertEqual(f('a:b:c'), ['a','b','c'])
212 def test_non_iterable(self):
213 """Verify a non-iterable string is caught"""
215 self.assertRaises(ArgumentTypeError, f, 0)
217 def test_type_conversion(self):
218 """Verify type conversion works properly"""
219 f = tolist(totype=int)
220 self.assertEqual(f('0,1,2'), [0, 1, 2])
222 ArgumentTypeError, f, '1,z,2,q')
224 def test_choices(self):
225 """Verify the choices validation works properly"""
226 f = tolist(totype=int, choices=[0, 1, 2, 3])
227 self.assertEqual(f('0,1,2'), [0, 1, 2])
228 self.assertRaises(ArgumentTypeError, f, '0,5,2')