简体   繁体   中英

Declare a static variable in an enum class

I have an enum for colors. I wish to add a helper method "toRGB()" to the enum class that converts an instance of the enum to an RGB object. As an optimization, I wished to create the dictionary once as a static variable. However, the correct syntax seems to elude me.

Can anyone suggest the right way to do this?

from enum import Enum

class RGB:
    def __init__(self, r, g, b):
        pass

class Color(Enum):
    RED = 0
    GREEN = 1

    __tbl = {
              RED:   RGB(1, 0, 0),
              GREEN: RGB(0, 1, 0)
            }

    def toRGB(self):
        return self.__class__.__tbl[self.value]

c = Color.RED
print(c.toRGB())

I get the following error:

Traceback (most recent call last):
  File "C:/Users/user/Desktop/test.py", line 20, in <module>
    print(c.toRGB())
  File "C:/Users/user/Desktop/test.py", line 17, in toRGB
    return self.__class__.__tbl[self.value]
TypeError: 'Color' object does not support indexing

As of Python 3.7, use the _ignore_ field: https://docs.python.org/3/library/enum.html

class Color(Enum):
    _ignore_ = ['_tbl']
    _tbl = {}  # nice for the type checker, but entirely ignored!


Color._tbl = {}  # actually creates the attribute

Non-method attributes become enum members (even tbl ). You can use a keyword argument instead:

class Color(Enum):
    RED = 0
    GREEN = 1

    def toRGB(self, tbl={
        RED:   RGB(1, 0, 0),
        GREEN: RGB(0, 1, 0)
    }):
        return tbl[self.value]

Alternatively, you can define the attribute after class creation:

class Color(Enum):
    RED = 0
    GREEN = 1

    def toRGB(self):
        return self._tbl[self]

Color._tbl = {
    Color.RED:   RGB(1, 0, 0),
    Color.GREEN: RGB(0, 1, 0)
}

We can't tell from your example if the 0 , 1 , 2 , ..., are meaningful values or just place-holders, but if they are just place-holders then the best solution is to discard them and use the RGB values as the Enum member value directly:

class Color(Enum):
    RED = 1, 0, 0
    GREEN = 0, 1, 0
    BLUE = 0, 0, 1

If the Enum members have a value separate from their rgb values you could use the new aenum library and solve the problem like this:

from aenum import Enum, NamedTuple

RGB = NamedTuple('RGB', 'r g b')

class Color(Enum, init='value rgb'):
    RED = 1, RGB(1,0,0)
    GREEN = 2, RGB(0,1,0)
    BLUE = 3, RGB(0,0,1)

and in use:

>>> Color.RED
<Color.RED: 1>

>>> Color.RED.rgb
RGB(r=1, g=0, b=0)

If what you want is a type conversion, then you could use the RGB class as the enumeration value:

from enum import Enum

class RGB:
    def __init__(self, r, g, b):
        # Check your inputs
        self.r = r
        self.g = g
        self.b = b

    def __str__(self):
        return f"{self.r} {self.g} {self.b}"

class Color(Enum):
    RED = RGB(1, 0, 0)
    GREEN = RGB(0, 1, 0)

    def toRGB():
        return c.value

c = Color.RED
print(c.toRGB())

You could also skip the toRGB helper method and simply write Color.RED.value

A different way to implement static class variables for an Enum subclass is to use the singleton pattern:

class Color(str, Enum):
    RED = "red"
    GREEN = "green"

    def __new__(cls, value):
        if not hasattr(cls, "instance"):
            cls.__tbl = {"red": RGB(1, 0, 0), "green": RGB(0, 1, 0)}
        obj = str.__new__(cls, value)
        return Enum.__new__(cls, obj)

Note that you cannot use instances of the enum like RED , GREEN etc. inside the definition of __new__ (they do not exist yet), so I have used string identifiers instead. I would prefer the first solution in general. My code looks like this since it needs to load a large dataset exactly once.

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM