diff --git a/hobbit_core/flask_hobbit/db.py b/hobbit_core/flask_hobbit/db.py index 496eb0b..f1282fb 100644 --- a/hobbit_core/flask_hobbit/db.py +++ b/hobbit_core/flask_hobbit/db.py @@ -1,5 +1,5 @@ # -*- encoding: utf-8 -*- -import enum +from enum import Enum, EnumMeta import six from sqlalchemy import Integer, Column, ForeignKey, func, DateTime @@ -61,7 +61,28 @@ def reference_col(tablename, nullable=False, pk_name='id', **kwargs): nullable=nullable, **kwargs) -class EnumExt(enum.Enum): +class EnumExtMeta(EnumMeta): + + def __new__(cls, name, bases, attrs): + obj = super(EnumExtMeta, cls).__new__(cls, name, bases, attrs) + + keys, values = set(), set() + for name, member in obj.__members__.items(): + member = member.value + if not isinstance(member, tuple) or len(member) != 2: + raise TypeError( + 'EnumExt member must be tuple type and length equal 2.') + key, value = member + if key in keys or value in values: + raise ValueError('duplicate values found: `{}`, please check ' + 'key or value.'.format(member)) + keys.add(key) + values.add(key) + + return obj + + +class EnumExt(six.with_metaclass(EnumExtMeta, Enum)): """ Extension for serialize/deserialize sqlalchemy enum field. Be sure ``type(key)`` is ``int`` and ``type(value)`` is ``str`` diff --git a/tests/test_db.py b/tests/test_db.py index 8764000..b9d4946 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -8,6 +8,19 @@ class TestEnumExt(BaseTest): + def test_key_index(self): + msg = "EnumExt member must be tuple type and length equal 2." + with pytest.raises(TypeError, message=msg): + class ErrTypeEnum(db.EnumExt): + CREATED = (0, u'新建', 'dda') + + msg = "ValueError: duplicate values found: `(0, '已完成')`, " + \ + "please check key or value." + with pytest.raises(ValueError, message=msg): + class ErrEnum(db.EnumExt): + CREATED = (0, u'新建') + FINISHED = (0, u'已完成') + @pytest.fixture def TaskState(self): class _TaskState(db.EnumExt): diff --git a/tox.ini b/tox.ini index ed926ee..e3e46ef 100644 --- a/tox.ini +++ b/tox.ini @@ -5,6 +5,7 @@ envlist = py27,py36,py37 deps = pytest pytest-cov + pytest-env flake8 py27: ipython==5.8.0 commands =