BitMaskEnum helper

gistfile1.py

from itertools import chain
from enum import Enum

class BitMaskEnum(object):
	
	def __init__(self, base_enum, iterable=None, value=None):
        if not isinstance(base_enum, Enum):
            raise TypeError('base_enum {!r} is not derived from {!r}'.format(base_enum, Enum))
		self._base = base_enum
		if iterable is not None:
			self._value = frozenset(base_enum(x) for x in iterable)
		elif value is not None:
			all_vals = [x.value for x in list(base_enum)]
			parts = []
			for mask in reversed(sorted(all_vals)):
				if mask == (mask & value):
					parts.append(base_enum(value))
			self._value = frozenset(parts)
		else:
			raise ValueError('Must provide on of iterable or value')

	def __iter__(self):
		for i in self._value:
			yield i

	@property
	def value(self):
		return sum(x.value for x in self._value)

	@property
	def name(self):
		return frozenset(x.name for x in self._value)

	def _set_operation(self, other, op):
		if not isinstance(other,  BitMaskEnum):
			if not hasattr(other, '__iter__'):
				other = BitMaskEnum(self._base, value=other)
			else:
				other = BitMaskEnum(self._base, iterable=other)
		return BitMaskEnum(self.base, iterable=op(self._value, other._value))


	def intersection(self, other):
		return _set_operation(self, other, frozenset.intersection)

	def difference(self, other):
		return _set_operation(self, other, frozenset.difference)

	def union(self, other):
		return _set_operation(self, other, frozenset.union)