CODEWITHSUNDEEP
pythonenumsbit-fields
dude
dude

Reputation: 53

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.

Upvotes: 2

Views: 268

Answers (2)

Ethan Furman
Ethan Furman

Reputation: 69031

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.

Upvotes: 3

0x0fba
0x0fba

Reputation: 1620

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

Upvotes: 3

Related Questions