Bug 1264482 - Add an enumeration-like string type with a limited set of possible values. r=ted

python/mozbuild/mozbuild/test/test_util.py

@@ -30,6 +30,8 @@ from mozbuild.util import (
     resolve_target_to_make,
     MozbuildDeletionError,
     HierarchicalStringList,
+    EnumString,
+    EnumStringComparisonError,
     ListWithAction,
     StrictOrderingOnAppendList,
     StrictOrderingOnAppendListWithFlagsFactory,
@@ -861,6 +863,29 @@ class TestMisc(unittest.TestCase):
             'before abc between a b c after'
         )
 
+class TestEnumString(unittest.TestCase):
+    def test_string(self):
+        CompilerType = EnumString.subclass('msvc', 'gcc', 'clang', 'clang-cl')
+
+        type = CompilerType('msvc')
+        self.assertEquals(type, 'msvc')
+        self.assertNotEquals(type, 'gcc')
+        self.assertNotEquals(type, 'clang')
+        self.assertNotEquals(type, 'clang-cl')
+        self.assertIn(type, ('msvc', 'clang-cl'))
+        self.assertNotIn(type, ('gcc', 'clang'))
+
+        with self.assertRaises(EnumStringComparisonError):
+            self.assertEquals(type, 'foo')
+
+        with self.assertRaises(EnumStringComparisonError):
+            self.assertNotEquals(type, 'foo')
+
+        with self.assertRaises(EnumStringComparisonError):
+            self.assertIn(type, ('foo', 'gcc'))
+
+        with self.assertRaises(ValueError):
+            type = CompilerType('foo')
 
 if __name__ == '__main__':
     main()

python/mozbuild/mozbuild/util.py

@@ -1137,3 +1137,38 @@ class DefinesAction(argparse.Action):
                 value = int(value)
         defines[name] = value
         setattr(namespace, self.dest, defines)
+
+
+class EnumStringComparisonError(Exception):
+    pass
+
+
+class EnumString(unicode):
+    '''A string type that only can have a limited set of values, similarly to
+    an Enum, and can only be compared against that set of values.
+
+    The class is meant to be subclassed, where the subclass defines
+    POSSIBLE_VALUES. The `subclass` method is a helper to create such
+    subclasses.
+    '''
+    POSSIBLE_VALUES = ()
+    def __init__(self, value):
+        if value not in self.POSSIBLE_VALUES:
+            raise ValueError("'%s' is not a valid value for %s"
+                             % (value, self.__class__.__name__))
+
+    def __eq__(self, other):
+        if other not in self.POSSIBLE_VALUES:
+            raise EnumStringComparisonError(
+                'Can only compare with %s'
+                % ', '.join("'%s'" % v for v in self.POSSIBLE_VALUES))
+        return super(EnumString, self).__eq__(other)
+
+    def __ne__(self, other):
+        return not (self == other)
+
+    @staticmethod
+    def subclass(*possible_values):
+        class EnumStringSubclass(EnumString):
+            POSSIBLE_VALUES = possible_values
+        return EnumStringSubclass