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.

as a side note, python has the `set` type for this kind of stuff. — gog, Dec 15 '22 at 11:29
My goal is to use the bitfield to store into pandas or numpy as uint32. — dude, Dec 15 '22 at 12:16

score 2 · Accepted Answer · edited Dec 15 '22 at 18:14

2

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

edited Dec 15 '22 at 18:14

Ethan Furman

63,992
20
159
237

answered Dec 15 '22 at 10:43

0x0fba

1,520
1
1
11

score 2 · Answer 2 · answered Dec 15 '22 at 18:57

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

@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.

Create an Enum with a bitfield

2 Answers2