Search code examples
pythonenumsbit-fields

Create an Enum with a bitfield

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.

Solution

  • 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)})