--- /dev/null
+
+"""
+Input handling routines
+Copyright 2014 Cray Inc. All Rights Reserved
+"""
+
+
+import itertools
+
+
+# Define some common integer multiplier suffixes
+
+# Powers of two
+binary_suffixes={
+ 'k': 2**10, 'K': 2**10,
+ 'm': 2**20, 'M': 2**20,
+ 'g': 2**30, 'G': 2**30,
+ 't': 2**40, 'T': 2**40,
+ 'p': 2**50, 'P': 2**50
+}
+memory_suffixes = binary_suffixes
+
+# Powers of ten
+decimal_suffixes={
+ 'k': 10**3, 'K': 10**3,
+ 'm': 10**6, 'M': 10**6,
+ 'g': 10**9, 'G': 10**9,
+ 't': 10**12, 'T': 10**12,
+ 'p': 10**15, 'P': 10**15
+}
+disk_suffixes = decimal_suffixes
+
+default_bases = [0, 16]
+
+def toint(string, base=default_bases, suffixes=binary_suffixes):
+ """Convert to integer with flexible base and multiplier support.
+
+ Provide a way to handle input that may be in any of several number
+ bases but may not use the appropriate prefix, e.g. 'deadbeef' rather
+ than the more pedantic '0xdeadbeef'. Also provide support for
+ multiplier suffixes, such as 'K' for kilo.
+
+ Arguments:
+
+ string - string to convert to integer
+ base - a single number, as used in int() or an iterable
+ of base values to try
+ suffixes - dictionary keyed by the string suffix with a value
+ to be used as a multiplier
+
+ The default base of [0, 16] allows the automatic recognition of numbers
+ with the standard prefixes and if that fails, tries a base 16 conversion.
+ """
+ try:
+ bases = list(base)
+ except TypeError:
+ # Object isn't iterable, so create one that is
+ bases = [base]
+
+ for b in bases:
+ if not (b == 0 or 2 <= b <= 36):
+ raise ValueError(
+ "toint() base {!s:s} must be >= 2 and <= 36".format(b))
+
+ multiplier = 1
+ try:
+ # Second iteration is after removing any suffix. This way, if
+ # a suffix happens to contain valid numeric characters, we'll
+ # try the numeric interpretation before we try their multiplier
+ # meaning, e.g. 'g' is a valid numeric value in base 17).
+ for i in xrange(2):
+ for b in bases:
+ try:
+ return int(string, b) * multiplier
+ except ValueError:
+ pass
+
+ if i != 0:
+ raise ValueError
+
+ # Find a suffix that matches the end of the string and use it
+ for k, v in suffixes.iteritems():
+ if string.endswith(k):
+ multiplier = v
+ string = string[0:-len(k)]
+ break
+ else:
+ raise ValueError
+
+ except ValueError:
+ suffix_list = suffixes.keys()
+ suffix_list.sort()
+ raise ValueError(
+ "invalid literal '{:s}' for toint() with base {!s:s} "
+ "and suffixes {!s:s}".format(string, list(bases), suffix_list))
+
+
+def hex2int(string):
+ """Wrapper for toint() which prefers base 16 input
+
+ This function is useful in situations where a callable must be passed,
+ such as with argparse.add_argument(type=hex2int, ...
+ """
+ return toint(string, base=[16, 0])
+
+
+def to_rangelist(args, default=xrange(0), base=[0,16],
+ suffixes=binary_suffixes):
+ """Convert a bunch of range list strings into a list of ranges
+
+ The arguments are:
+
+ args - iterable containing ranglist strings
+ default - iterator to return if args is empty
+ base - number base to use for integer conversion
+ suffixes - integer multiplier suffixes
+
+ Each arg is taken to be a range list, where a range list may be:
+
+ rangelist ::= range[,range]...
+ range ::= <first>-<last> | <first>#<count> | <value>
+
+ where the range first-last is inclusive.
+ """
+ if len(args) == 0:
+ return default
+
+ ranges = []
+ for range_list_str in args:
+ range_strs = range_list_str.split(',')
+ for range_str in range_strs:
+ if "-" in range_str:
+ fields = range_str.split('-', 1)
+ start = toint(fields[0], base, suffixes=suffixes)
+ end = toint(fields[1], base, suffixes=suffixes) + 1
+ ranges.append(xrange(start, end))
+ elif "#" in range_str:
+ fields = range_str.split('#', 1)
+ start = toint(fields[0], base, suffixes=suffixes)
+ end = start + toint(fields[1], base, suffixes=suffixes)
+ ranges.append(xrange(start, end))
+ else:
+ start = toint(range_str, base, suffixes=suffixes)
+ end = start + 1
+ ranges.append(xrange(start, end))
+
+ return ranges
+
+
+def iter_rangestr(*args, **kwargs):
+ """Convert a bunch of range list strings into a single iterator
+
+ The arguments are the same as for to_rangelist().
+ """
+ return itertools.chain(*to_rangelist(*args, **kwargs))
+
+
+if __name__ == '__main__':
+ import unittest
+
+ # toint()
+ class Test_toint(unittest.TestCase):
+ def test_base_zero(self):
+ self.assertEqual(toint('0b10', 0), 2)
+ self.assertEqual(toint('0o10', 0), 8)
+ self.assertEqual(toint('10', 0), 10)
+ self.assertEqual(toint('0x10', 0), 16)
+
+ def test_base_out_of_range(self):
+ self.assertRaises(ValueError, toint, '10', -1)
+ self.assertRaises(ValueError, toint, '10', 1)
+ self.assertRaises(ValueError, toint, '10', 37)
+
+ def test_base_search(self):
+ bases = [0, 16]
+ self.assertEqual(toint('10', bases), 10)
+ self.assertEqual(toint('f', bases), 15)
+
+ self.assertEqual(toint('0b10', bases), 2)
+ self.assertEqual(toint('0o10', bases), 8)
+ self.assertEqual(toint('10', bases), 10)
+ self.assertEqual(toint('0x10', bases), 16)
+
+ def test_suffixes(self):
+ for k, v in binary_suffixes.iteritems():
+ self.assertEqual(toint('0b10'+k), 0b10*v)
+ self.assertEqual(toint('0o10'+k), 0o10*v)
+ self.assertEqual(toint('10'+k), 10*v)
+ self.assertEqual(toint('0x10'+k), 0x10*v)
+
+ def test_suffix_number_overlap(self):
+ # Verify a valid numeric isn't used as a suffix
+ self.assertEqual(toint('1g', 17), 33)
+ self.assertEqual(toint('1gk', 17), 33*binary_suffixes['k'])
+
+
+ # hex2int()
+ class Test_hex2int(unittest.TestCase):
+ """Verify the hex2int() function"""
+ def test_explicit_base(self):
+ """Verify that explicit base syntax is honored"""
+ self.assertEqual(hex2int('0x10'), 16)
+ self.assertEqual(hex2int('0o10'), 8)
+
+ def test_default_base(self):
+ """Verify that base 16 is preferred"""
+ self.assertEqual(hex2int('10'), 16)
+ self.assertEqual(hex2int('0b10'), 2832)
+
+
+ # iter_rangelist()
+ class Test_iter_rangelist(unittest.TestCase):
+ """Test both iter_rangelist and the underlying to_rangelist."""
+ def test_good_single_ranges(self):
+ self.assertEqual(list(iter_rangestr([])), [])
+ self.assertEqual(list(iter_rangestr(['1-2'])), list(xrange(1,3)))
+ self.assertEqual(list(iter_rangestr(['1#2'])), list(xrange(1,3)))
+ self.assertEqual(list(iter_rangestr(['1'])), list(xrange(1,2)))
+
+ def test_good_multiple_ranges(self):
+ test_rangestrs = [
+ # Test params, Expected result
+ (['1', '3-5', '1#2'], [1, 3, 4, 5, 1, 2]),
+ ]
+
+ for ranges, expected in test_rangestrs:
+ # Test the ranges as separate list elements
+ self.assertEqual(list(iter_rangestr(ranges)), expected)
+
+ # Test the ranges joined by commas
+ joined = [','.join(ranges)]
+ self.assertEqual(list(iter_rangestr(joined)), expected)
+
+ def test_bad_single_ranges(self):
+ self.assertRaises(ValueError, iter_rangestr, ['1#2#3'])
+ self.assertRaises(ValueError, iter_rangestr, ['1#2-3'])
+ self.assertRaises(ValueError, iter_rangestr, ['1-2#3'])
+ self.assertRaises(ValueError, iter_rangestr, ['1-2-3'])
+
+ # Run all unit tests
+ unittest.main()