Search code examples
pythonenumscomplextype

Python enums with complex types


I'm new to Python and I'm wondering if I can build enums with complex structures, not just primitive types. For instance (in pseudo-code):

Point::Enum
  x, y
  constructor ( x, y ) {
    ...
  }

  bottom_left = Point ( 0, 0 )
  top_left = Point ( 0, 100 )
  top_right = Point ( 100, 100 )
  bottom_right = Point ( 100, 0 )

So far, I could only find Python documentation that mentions enums with strings or ints.


Solution

  • If you want Point as a separate entity from the Enum that tracks the corners, then you need them to be separate:

    from enum import Enum
    
    class Point(object):
        def __init__(self, x, y):
            self.x = x
            self.y = y
        def __repr__(self):
            return 'Point(%r, %r)' % (self.x, self.y)
    
    class Corner(Enum):
        BottomLeft = Point(0, 0)
        TopLeft = Point(0, 100)
        TopRight = Point(100, 100)
        BottmRight = Point(100, 0)
    

    Doing it this way means that each enum contains a Point as its value, but is not a Point itself:

    >>> Corner.BottomLeft
    <Corner.BottomLeft: Point(0, 0)>
    >>> Corner.BottomLeft.value
    Point(0, 0)
    

    If you want the enum members to be a Point, then mix in the Point class:

    from enum import Enum
    
    class Point(object):
        def __init__(self, x, y):
            self.x = x
            self.y = y
        def __repr__(self):
            return 'Point(%r, %r)' % (self.x, self.y)
    
    class Corner(Point, Enum):
        BottomLeft = 0, 0
        TopLeft = 0, 100
        TopRight = 100, 100
        BottmRight = 100, 0
    
    >>> Corner.TopLeft
    <Corner.TopLeft: (0, 0)>
    >>> isinstance(Corner.TopLeft, Point)
    True
    >>> Corner.TopLeft.value
    (0, 100)
    >>> Corner.TopLeft.x
    0
    >>> Corner.TopLeft.y
    100
    

    Finally, if all you need is for the enums to have the x and y attributes:

    from aenum import Enum
    
    class Corner(Enum):
        __init__ = 'x y'
        BottomLeft = 0, 0
        TopLeft = 0, 100
        TopRight = 100, 100
        BottmRight = 100, 0
    
    >>> Corner.TopLeft
    <Corner.TopLeft: (0, 100)>
    >>> Corner.TopLeft.value
    (0, 100)
    >>> Corner.TopLeft.x
    0
    >>> Corner.TopLeft.y
    100
    

    Note that that last example is using the aenum package1. You can accomplish the same thing with either enum34 or the stdlib enum by writing an __init__ for the Point class.


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