Commits

Felix Krull committed 4ca8ad6

Add Enum type.

Comments (0)

Files changed (2)

provence/core/types.py

 
     def __eq__(self, other):
         return isinstance(other, Seq) and self.type == other.type
+
+
+class Enum(AttributeType):
+    def __init__(self, *values):
+        self.values = frozenset(values)
+
+    def check_type(self, value):
+        return value in self.values
+
+    def __str__(self):
+        return "Enum(%s)" % ", ".join(self.values)
+
+    def __eq__(self, other):
+        return isinstance(other, Enum) and self.values == other.values

provence/tests/test_types.py

 import unittest
-from provence.core.types import check_type, Maybe, OneOf, Seq, Number
+from provence.core.types import check_type, Enum, Maybe, OneOf, Seq, Number
 
 
 class CheckTypeTestCase:
     def test_invalid_types(self):
         with self.assertRaises(TypeError):
             check_type(5, "not a type")
+
+
+class TestEnum(CheckTypeTestCase, unittest.TestCase):
+    CHECK_TRUE = [
+        ("yes", Enum("yes", "no")),
+        ("no", Enum("yes", "no")),
+        (5, Enum(*range(6))),
+        ("str", Enum("str")),
+    ]
+
+    CHECK_FALSE = [
+        ("NO", Enum("yes", "no")),
+        (6, Enum(*range(6))),
+        ("no", Enum("str")),
+    ]
+
+    def test_eq(self):
+        self.assertEqual(Enum("yes", "no"), Enum("yes", "no"))
+        self.assertEqual(Enum("no", "yes"), Enum("yes", "no"))
+        self.assertNotEqual(Enum(1), Enum(2))
+        self.assertEqual(Enum("yes", "no", "yes", "no"),
+                         Enum("no", "yes", "yes"))