enum masks idiom in python

We've all seen this:

class Foo {
public:
  enum FooType
  {
    FooTypeOne    = 0x1,
    FooTypeTwo   = 0x1 << 1,
    FooTypeThree = 0x1 << 2,
    FLAGS = ONE | TWO | THREE | FOUR
  };
}

I needed to implement something similar in python. Before doing so, I check to see what the canonical way of doing so is (e.g. 'switch' vs. 'dispatch dictionary'). Here it is:

ONE, TWO, THREE, FOUR = [frozenset([x]) for x in 'one two three four'.split()]
FLAGS = frozenset([ONE, TWO, THREE, FOUR])

class Foo(object):
    def print_flag(self, mask):
        for flag in FLAGS:
            if flag & mask:
                print(flag)

foo = Foo()

foo.print_flag(ONE | TWO)

Output:

frozenset(['one'])
frozenset(['two'])

In addition to the expected binary masking functionality by leveraging the bitwise math operators of sets, we get a human readable text "label" for printing or displaying the flag. (thanks to Aaron Gallagher @ #python)

Share on: TwitterFacebookGoogle+Email

Comments !