Search code examples
pythonenumspython-typing

Type hinting an enum property that returns an instance of that enum


I have an enum that looks like this:

class Direction(Enum):
    NORTH = 0
    EAST = 1
    SOUTH = 2
    WEST = 3

    @property
    def left(self) -> Direction:
        new_direction = (self.value - 1) % 4
        return Direction(new_direction)

    @property
    def right(self) -> Direction:
        new_direction = (self.value + 1) % 4
        return Direction(new_direction)

I am trying to type hint the left and right properties to indicate their return value is of type Direction.

I thought the above would work, but when I run the code I get the following error: NameError: name 'Direction' is not defined. I guess this is happening because the Python interpreter does not yet know what the Direction enum is at the time of defining this function.

My question is, is there anyway I can type hint these properties? Thanks.


Solution

  • This is known as a forward reference as the Direction class isn't defined yet when the property function signature is executed. You'll need to put the forward reference in quotes. See https://www.python.org/dev/peps/pep-0484/#forward-references for more information.

    from enum import Enum
    
    class Direction(Enum):
        NORTH = 0
        EAST = 1
        SOUTH = 2
        WEST = 3
    
        @property
        def left(self) -> 'Direction':
            new_direction = (self.value - 1) % 4
            return Direction(new_direction)
    
        @property
        def right(self) -> 'Direction':
            new_direction = (self.value + 1) % 4
            return Direction(new_direction)