How to add some data to an Enum in Python

Question

I have a class that I use to define different types of plots I am performing

class MyPlots(Enum):
 STANDARDSCALE = "standard"
 LOGSCALE = "log"

there are default values associated with the coordinates I usually plot with, i.e., in my code I have parts like

 ...
 if self._plottype==MyPlots.STANDARDSCALE:
 xlimits = [math.e, math.e**2]
 if self._plottype==MyPlots.LOGSCALE:
 xlimits = [1,2]
 ...

I am wondering about what would be the best way to refactor this code, i.e., writing an abstract class

class PlotTypeWorks(ABC):
 xlimits
 @virutalmethod
 def __init__(self, type : MyPlots):
 """ Sets limits """
class Plot(PlotTypeWorks):
 def __init__(self, type : MyPlots):
 self._plottype = type
 if self._plottype==MyPlots.STANDARDSCALE:
 self.xlimits = [math.e, math.e**2]
 if self._plottype==MyPlots.LOGSCALE:
 self.xlimits = [1,2]

or having PlotStandard and PlotLog classes.

I want the code to be fast and easy to read at the same time. What I am concerned about, is that with different classes (PlotStandard and PlotLog) the code might soon be overcrowded with plots with different parameters (hard to read). From the other side, with only one class, I might be calling a lot of checks (if self._plottype==X) that might just slow my code when I will create thousands of those classes.

Is there a solution with the perks of both? i.e., less checks and more readability?

• ThePunisher: I fail to see how the check would be the bottleneck that slows your code.

Answer 1

There are a couple different ways you can build that information directly into the enum definitions:

1. use the limit values instead of the string for the enum value;
2. add the limit values to the string value

(1) would look like:

class MyPlots(Enum):
 STANDARDSCALE = math.e, math.e**2
 LOGSCALE = 1, 2
xlimits = self._plottype.value

(2) would look like

class MyPlots(Enum):
 STANDARDSCALE = "standard", (math.e, math.e**2)
 LOGSCALE = "log", (1, 2)
 #
 def __new__(cls, value, limits):
 member = object.__new__(cls)
 member._value_ = value
 member.limits = limits
 return member
xlimits = self._plottype.limits

In use, (2) would look like:

>>> list(MyPlots)
[<MyPlots.STANDARDSCALE: 'standard'>, <MyPlots.LOGSCALE: 'log'>]
>>> MyPlots.LOGSCALE.limits
(1, 2)

• La Cartuccia: Can you provide an example on how you would then create an object of the second kind? It is very interesting but it looks like you should provide value and limits in the initializer (is __new__ the same as __init__?)
• Ethan Furman: @LaCartuccia: No, __new__ creates and returns an object while __init__ works with the already created object. Typically, mutable objects do their setup work in __init__ and immutable objects do it in __new__. It is possible to split the work and have both a __new__ and an __init__, but in this case that doesn't provide any benefit.
• Ethan Furman: @LaCartuccia: Not sure what you mean by "an object of the second kind". I added an example use of my number (2) -- hopefully that helps.
• La Cartuccia: Oh I see thank you. And yes, I meant number (2). So the idea is that limits is now a static member, right? I had a look at the documentation and they have an example for Planet
• Ethan Furman: @LaCartuccia: It depends on the goal -- in my (2) if the __new__ wasn't there then the value would be the entire assignment "log", (1,2).