Skip to content

Commit 24a4810

Browse files
committed
add implementation of IntFlag as ListField of IntEnumField
1 parent 5e63e05 commit 24a4810

File tree

1 file changed

+36
-1
lines changed

1 file changed

+36
-1
lines changed

extras_mongoengine/fields.py

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
1+
import operator
2+
13
from datetime import timedelta
24
from mongoengine.base import BaseField
3-
from mongoengine.fields import IntField, StringField, EmailField
5+
from mongoengine.fields import EmailField, IntField, ListField, StringField
6+
7+
try:
8+
from functools import reduce
9+
except ImportError:
10+
# reduce is a builtin in Python2
11+
pass
412

513

614
class TimedeltaField(BaseField):
@@ -105,6 +113,33 @@ class IntEnumField(EnumField, IntField):
105113
pass
106114

107115

116+
class IntFlagField(ListField):
117+
def __init__(self, enum, **kwargs):
118+
super(IntFlagField, self).__init__(IntEnumField(enum), **kwargs)
119+
120+
def __get__(self, instance, owner):
121+
if instance is None:
122+
return self
123+
124+
return self.field.enum(reduce(
125+
operator.or_, instance._data.get(self.name, []), 0))
126+
127+
def __set__(self, instance, value):
128+
# copy mongoengine
129+
if value is None:
130+
if self.null:
131+
value = None
132+
elif self.default is not None:
133+
value = self.default
134+
if callable(value):
135+
value = value()
136+
137+
if value is not None and not isinstance(value, list):
138+
value = [i for i in self.field.enum if i and i & value == i]
139+
140+
super(IntFlagField, self).__set__(instance, value)
141+
142+
108143
class StringEnumField(EnumField, StringField):
109144
"""A variation on :class:`EnumField` for only string containing enumeration.
110145
"""

0 commit comments

Comments
 (0)