Modify semantics of UnitsParser's convertible_to parameter.

perfkitbenchmarker/flag_util.py

@@ -227,26 +227,30 @@ def __exit__(self, *unused_args, **unused_kwargs):
 class UnitsParser(flags.ArgumentParser):
   """Parse a flag containing a unit expression.
 
-  The user may require that the provided expression is convertible to
-  a particular unit. The parser will throw an error if the condition
-  is not satisfied. For instance, if a unit parser requires that its
-  arguments are convertible to bits, then KiB and GB are valid units
-  to input, but meters are not. If the user does not require this,
-  than *any* unit expression is allowed.
+  Attributes:
+    convertible_to: list of units.Unit instances. A parsed expression must be
+        convertible to at least one of the Units in this list. For example,
+        if the parser requires that its inputs are convertible to bits, then
+        values expressed in KiB and GB are valid, but values expressed in meters
+        are not.
   """
 
   syntactic_help = ('A quantity with a unit. Ex: 12.3MB.')
 
-  def __init__(self, convertible_to=None):
+  def __init__(self, convertible_to):
     """Initialize the UnitsParser.
 
     Args:
-      convertible_to: units.Unit or
-        None. If a unit, the input must be convertible to this unit or
-        the Parse() method will raise a ValueError.
+      convertible_to: Either an individual unit specification or a series of
+          unit specifications, where each unit specification is either a string
+          (e.g. 'byte') or a units.Unit. The parser input must be convertible to
+          at least one of the specified Units, or the Parse() method will raise
+          a ValueError.
     """
-
-    self.convertible_to = convertible_to
+    if isinstance(convertible_to, (basestring, units.Unit)):
+      self.convertible_to = [units.Unit(convertible_to)]
+    else:
+      self.convertible_to = [units.Unit(u) for u in convertible_to]
 
   def Parse(self, inp):
     """Parse the input.
@@ -259,9 +263,9 @@ def Parse(self, inp):
       A units.Quantity.
 
     Raises:
-      ValueError if it can't parse its input.
+      ValueError: If the input cannot be parsed, or if it parses to a value with
+          improper units.
     """
-
     if isinstance(inp, units.Quantity):
       quantity = inp
     else:
@@ -273,12 +277,16 @@ def Parse(self, inp):
       if not isinstance(quantity, units.Quantity):
         raise ValueError('Expression %r evaluates to a unitless value.' % inp)
 
-    if self.convertible_to is not None:
+    for unit in self.convertible_to:
       try:
-        quantity.to(self.convertible_to)
+        quantity.to(unit)
+        break
       except units.DimensionalityError:
-        raise ValueError("Expression %s is not convertible to %s" %
-                         (inp, self.convertible_to))
+        pass
+    else:
+      raise ValueError(
+          'Expression {0!r} is not convertible to an acceptable unit '
+          '({1}).'.format(inp, ', '.join(str(u) for u in self.convertible_to)))
 
     return quantity
 
@@ -288,22 +296,22 @@ def Serialize(self, units):
     return str(units)
 
 
-def DEFINE_units(name, default, help, convertible_to=None,
+def DEFINE_units(name, default, help, convertible_to,
                  flag_values=flags.FLAGS, **kwargs):
   """Register a flag whose value is a units expression.
 
   Args:
     name: string. The name of the flag.
     default: units.Quantity. The default value.
     help: string. A help message for the user.
-    convertible_to: units.Unit or None. If a unit is provided, the input must be
-        convertible to this unit to be considered valid.
+    convertible_to: Either an individual unit specification or a series of unit
+        specifications, where each unit specification is either a string (e.g.
+        'byte') or a units.Unit. The flag value must be convertible to at least
+        one of the specified Units to be considered valid.
     flag_values: the gflags.FlagValues object to define the flag in.
   """
-
   parser = UnitsParser(convertible_to=convertible_to)
   serializer = UnitsSerializer()
-
   flags.DEFINE(parser, name, default, help, flag_values, serializer, **kwargs)
 
 

tests/flag_util_test.py

@@ -134,8 +134,9 @@ def testReadAndWrite(self):
 
 
 class TestUnitsParser(unittest.TestCase):
+
   def setUp(self):
-    self.up = flag_util.UnitsParser()
+    self.up = flag_util.UnitsParser('byte')
 
   def testParser(self):
     self.assertEqual(self.up.Parse('10KiB'), 10 * 1024 * units.byte)
@@ -162,21 +163,22 @@ def testBytesWithPrefix(self):
     self.assertEqual(q.magnitude, 2000.0)
     self.assertEqual(q.units, {'byte': 1.0})
 
-  def testPercent(self):
-    q = self.up.Parse('10%')
-    self.assertEqual(q.magnitude, 10)
-    self.assertEqual(q.units, units.percent)
+  def testWrongUnit(self):
+    with self.assertRaises(ValueError):
+      self.up.Parse('1m')
 
-  def testConvertibleTo(self):
+  def testConvertibleToUnit(self):
     up = flag_util.UnitsParser(convertible_to=units.byte)
     self.assertEqual(up.Parse('10KiB'), 10 * 1024 * units.byte)
 
-  def testConvertibleToWrongUnit(self):
-    up = flag_util.UnitsParser(convertible_to=units.byte)
+  def testConvertibleToSeries(self):
+    up = flag_util.UnitsParser(convertible_to=(units.byte, 'second'))
+    self.assertEqual(up.Parse('10 MB'), 10 * units.Unit('megabyte'))
+    self.assertEqual(up.Parse('10 minutes'), 10 * units.Unit('minute'))
     with self.assertRaises(ValueError):
-      up.Parse('1m')
+      up.Parse('1 meter')
 
-  def testConvertibleToPercent(self):
+  def testPercent(self):
     up = flag_util.UnitsParser(convertible_to=units.percent)
     self.assertEqual(up.Parse('100%'), 100 * units.percent)
     with self.assertRaises(ValueError):