Create an Enum with a bitfield

Question

First a bit of context: I want to process various documents and detect for each one which european countries have been mentioned.

Example:

EU_COUNTRIES = (
    "Austria",
    "Belgium",
    "Bulgaria",
    "Croatia"
)  # 4 out of the 27

To optimize ram consumption (and avoid 27 columns) I'm going to use a bitfield. For a document where Austria and Bulgaria are mentioned, the bitfield equals to 5.

I'd like to have an enum to help testing against the bitfield.

For example: is Bulgaria mentioned in this document?

What I expect: 5 & Country.Bulgaria == True .

What I've tried:

from enum import Enum

class Country(Enum):
    Austria = 1
    Belgium = 2
    Bulgaria = 4
    Croatia = 8

print(5 & Country.Bulgaria)

I've created this enum but got error: TypeError: unsupported operand type(s) for &: 'int' and 'Country'

How can I achieve the goal (enum or other solution).

Thanks for your help.

Answer 1

It will work with an IntEnum : "Enum where members are also (and must be) ints".

Additionally you can build your enum by using the functional syntax . Combined with bitshift operations it allows to easily generate the enum.

from enum import IntEnum

Country = IntEnum("Country", {c:1<<i for i, c in enumerate(EU_COUNTRIES)})

Answer 2

Flag was created for this purpose (or IntFlag if you need to work directly with int s).

@0x0fba's usage of the Functional API is excellent; if your countries are in a separate file you might want to adapt this answer instead to create the flag.

Either way, once you have a Country class you would use it something like:

affected_countries = Country(5)  # `5` would actually be taken from the data
if Country.Bulgaria in affected_countries:
    # do stuff

---

Disclosure: I am the author of the Python stdlib Enum , the enum34 backport , and the Advanced Enumeration ( aenum ) library.