utils: add Enum method.
BUG=none
TEST=make test; make lint
Change-Id: I87e7cddf54d8d2b1b27b3dae9404a99a95d66811
diff --git a/graphyte/utils.py b/graphyte/utils.py
index 732a422..f971d3a 100644
--- a/graphyte/utils.py
+++ b/graphyte/utils.py
@@ -6,8 +6,10 @@
"""Utils class and functions. """
+import collections
from contextlib import contextmanager
import inspect
+import itertools
import math
import os
import tempfile
@@ -18,6 +20,22 @@
from .default_setting import logger
+def Enum(*args, **kwargs):
+ """Create the immutable enumeration set.
+
+ Usage:
+ 1. C-style enum. The value starts from 0 and increase orderly.
+ A = Enum('foo', 'bar')
+ then A.foo == 0, A.bar == 1
+ 2. Key-value pair enum.
+ B = Enum(foo='FOO', bar='BAR')
+ then B.foo == 'FOO', B.bar == 'BAR'
+ """
+ fields = dict(zip(args, itertools.count()), **kwargs)
+ enum_type = collections.namedtuple('Enum', fields.keys())
+ return enum_type(**fields)
+
+
def SearchConfig(filepath, search_dirs=None):
"""Find the config file and return the content.
The order of searching is:
diff --git a/graphyte/utils_unittest.py b/graphyte/utils_unittest.py
index dd459ab..19b3d77 100755
--- a/graphyte/utils_unittest.py
+++ b/graphyte/utils_unittest.py
@@ -106,3 +106,30 @@
now = time.time()
with self.assertRaises(utils.TimeoutError):
utils.WaitFor(lambda: _ReturnTrueAfter(now + 1), timeout_secs=0.5)
+
+
+class EnumTest(unittest.TestCase):
+ def testArgsEnum(self):
+ ret = utils.Enum('A', 'B', 'C')
+ self.assertEquals(ret.A, 0)
+ self.assertEquals(ret.B, 1)
+ self.assertEquals(ret.C, 2)
+
+ def testKwargsEnum(self):
+ ret = utils.Enum(A='a', B=1, C=False)
+ self.assertEquals(ret.A, 'a')
+ self.assertEquals(ret.B, 1)
+ self.assertEquals(ret.C, False)
+
+ def testArgsAndKwargsEnum(self):
+ ret = utils.Enum('foo', 'bar', A='a', B=1, C=None)
+ self.assertEquals(ret.foo, 0)
+ self.assertEquals(ret.bar, 1)
+ self.assertEquals(ret.A, 'a')
+ self.assertEquals(ret.B, 1)
+ self.assertEquals(ret.C, None)
+
+ def testEnumImmutable(self):
+ ret = utils.Enum('foo')
+ with self.assertRaises(AttributeError):
+ ret.foo = 'NO WAY!!'
