Whamcloud - gitweb
LU-12461 contrib: Add epython scripts for crash dump analysis
[fs/lustre-release.git] / contrib / debug_tools / epython_scripts / crashlib / input / argparse_ext.py
1
2 """
3 Module which provides extensions for the standard Python argparse module.
4 Copyright 2015 Cray Inc.  All Rights Reserved
5 """
6
7 import argparse
8 import copy
9
10 from argparse import _ensure_value, Action, ArgumentTypeError
11
12
13 class ExtendAction(Action):
14     """Action to extend a list of argument values
15
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:
19
20     -   This action must receive an iterable 'values' argument from the
21         parser.  There are two ways to make this happen:
22
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
26             is not '?'
27     """
28
29     def __call__(self, parser, namespace, values, option_string):
30         items = copy.copy(_ensure_value(namespace, self.dest, []))
31
32         try:
33             items.extend(values)
34         except TypeError:
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__))
39
40         setattr(namespace, self.dest, items)
41
42
43 def str2list(string, sep=',', totype=None, choices=None):
44     """Split a string into a list with conversion and validation
45
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.
49     """
50
51     plural = {False: '', True: 's'}
52
53     # Values should be string or an iterable container of strings.
54     # Split values on the separator into a list
55     try:
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__))
61
62     # Perform type conversion
63     if totype is not None:
64         errs = []
65         for i, v in enumerate(lst):
66             try:
67                 lst[i] = totype(v)
68             except (TypeError, ValueError):
69                 errs.append(v)
70         if errs:
71             msg = "invalid {:s} value{:s}: {!r:s}".format(
72                 totype.__name__, plural[len(errs) > 1], errs)
73             raise ArgumentTypeError(msg)
74
75     # Verify each separate value
76     if choices is not None:
77         errs = filter(lambda x:x not in choices, lst)
78         if errs:
79             msg = "invalid choice{:s}: {!r:s} (choose from {!s:s})".format(
80                 plural[len(errs) > 1], errs, choices)
81             raise ArgumentTypeError(msg)
82
83     return lst
84
85
86 def tolist(sep=',', totype=None, choices=None):
87     """Returns a parameterized callable for argument parser type conversion
88
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
91     behavior.
92     """
93     return lambda x:str2list(x, sep=sep, totype=totype, choices=choices)
94
95
96 if __name__ == '__main__':
97     import unittest
98
99     class Test_Action_Base(unittest.TestCase):
100         """Create a base class for testing argparse Action classes"""
101
102         def setUp(self):
103             """Create the ExtendAction object and args Namespace"""
104             self.action = ExtendAction([], dest='dest')
105             self.args   = argparse.Namespace()
106
107         def actionRun(self, values):
108             """Run the Action instance using values"""
109             self.action(None, self.args, values, '')
110
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)
115
116         def actionArgTypeErr(self, values):
117             """Run the Action and verify it raises ArgumentTypeError"""
118             self.assertRaises(
119                 ArgumentTypeError, self.action, None, self.args, values, '')
120
121
122     class Test_ExtendAction(Test_Action_Base):
123         """Test the ExtendAction class"""
124
125         def test_non_iterable(self):
126             """Test ExtendAction with a non-iterable type
127
128             This is similar to:
129                 parser.add_argument('-z', nargs=None, type=int ...)
130
131                 parser.parse_args(['-z', '0'])
132             """
133             self.actionArgTypeErr(0)
134
135         def test_single_value(self):
136             """Test ExtendAction with a single value
137
138             This is similar to:
139                 parser.add_argument('-z', nargs=None ...)
140
141                 parser.parse_args(['-z', 'a'])
142             """
143             self.actionEqual('a', ['a'])
144
145         def test_single_string(self):
146             """Test ExtendAction with a single value
147
148             This is similar to:
149                 parser.add_argument('-z', nargs=None ...)
150
151                 parser.parse_args(['-z', 'abc'])
152             """
153             self.actionEqual('abc', ['a', 'b', 'c'])
154
155         def test_single_value_multiple_calls(self):
156             """Test ExtendAction with a single value and multiple calls
157
158             This is similar to:
159                 parser.add_argument('-z', nargs=None, type=int ...)
160
161                 parser.parse_args(['-z', 'a', '-z', 'b'])
162             """
163             self.actionEqual('a', ['a'])
164             self.actionEqual('b', ['a', 'b'])
165
166         def test_value_list(self):
167             """Test ExtendAction with a value list
168
169             This is similar to:
170                 parser.add_argument('-z', nargs=1 ...)
171
172                 parser.parse_args(['-z', 'abc'])
173             """
174             self.actionEqual(['abc'], ['abc'])
175
176         def test_value_list_multiple_calls(self):
177             """Test ExtendAction with a single value and multiple calls
178
179             This is similar to:
180                 parser.add_argument('-z', nargs=1 ...)
181
182                 parser.parse_args(['-z', 'abc', '-z', 'def'])
183             """
184             self.actionRun(['abc'])
185             self.actionEqual(['def'], ['abc', 'def'])
186
187         def test_value_list_multiple_values(self):
188             """Test ExtendAction with a value list of length > 1
189
190             This is similar to:
191                 parser.add_argument('-z', nargs=2 ...)
192                 -or-
193                 parser.add_argument('-z', nargs='+' ...)
194                 -or-
195                 parser.add_argument('-z', nargs='*' ...)
196
197                 parser.parse_args(['-z', 'abc', 'def'])
198             """
199             self.actionEqual(['abc', 'def'], ['abc', 'def'])
200
201
202     class Test_tolist_str2list(unittest.TestCase):
203         """Test the str2list and tolist conversion functions"""
204
205         def test_sep(self):
206             """Verify default and non-default separators work"""
207             f = tolist()
208             self.assertEqual(f('a,b,c'), ['a','b','c'])
209             f = tolist(sep=':')
210             self.assertEqual(f('a:b:c'), ['a','b','c'])
211
212         def test_non_iterable(self):
213             """Verify a non-iterable string is caught"""
214             f = tolist()
215             self.assertRaises(ArgumentTypeError, f, 0)
216
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])
221             self.assertRaises(
222                 ArgumentTypeError, f, '1,z,2,q')
223
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')
229
230
231     unittest.main()