Can I decompose a Python Flag enum into its fundamental parts?

Can I decompose a Python Flag enum into its fundamental parts?

Problem Description:

I have an enum that represents the directions you’re allowed to move for a given cell in a maze:

class Direction(Flag):
    NORTH = 1
    EAST = 2
    SOUTH = 4
    WEST = 8
    NE = NORTH | EAST
    NW = NORTH | WEST
    ...etc
    NESW = NORTH | EAST | SOUTH | WEST

This makes it easy to check if you can go west, you can just check cell.directions & Direction.WEST. But what if I want to iterate over the possible directions? Something like for d in cell.directions: ..., but you can’t do this. If I didn’t alias every possible combination of directions (which makes things easier in my code), then I could do:

for d in Direction:
    if cell.directions & d:
        ....

But this won’t work for me because it would iterate over all the combined directions too, rather than just the four basic cardinal directions. Is there a good solution here?

Solution – 1

In Python 3.11 this works as you would expect: iteration only provides the canonical (i.e. powers of two) flags, all others are considered aliases.

Even in prior versions, though, this works:

from enum import Flag

class Direction(Flag):
    NORTH = 1
    EAST = 2
    SOUTH = 4
    WEST = 8

NORTH, EAST, SOUTH, WEST = Direction

directions = NORTH | EAST
d = NORTH
if d in directions:
    print('going North!')

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

Solution – 2

You could just create a list basic_directions = [NORTH, EAST, SOUTH, WEST] and use the same for loop that you wrote before:

for d in basic_directions:
    if cell.directions & d:
        ....
Rate this post
We use cookies in order to give you the best possible experience on our website. By continuing to use this site, you agree to our use of cookies.
Accept
Reject