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?

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)