CODEWITHSUNDEEP
pythonpython-3.xenums
Dragan R.
Dragan R.

Reputation: 670

Enum function to map value to a related value

Currently, I have an enum that represents a name/number mapping. Each EventType, however, also has additional "related properties" (e.g. a status code, and a message type).

class EventType(Enum):
  CANCELLED = 0
  ONTIME = 1
  DELAYED = 2

  def get_status(self):
    if self == EventType.CANCELLED:
        return "statuscode1"
    elif self == EventType.DELAYED:
        return "statuscode2"
    else:
        return "statuscode3"

  def get_message_type(self):
    if self == EventType.CANCELLED:
        return "messagetype1"
    elif self == EventType.DELAYED:
        return "messagetype2"
    else:
        return "messagetype3"

Instead of creating the methods above and tons of if chains checking against self, is there a cleaner way of refactoring to return the status codes and message types? It's almost as if CANCELLED = (0, statuscode1, messagetype1).. How can I represent this concept in an enum? Is an enum even the correct way to do this?

Upvotes: 5

Views: 8457

Answers (3)

Merelda
Merelda

Reputation: 1338

In Python 3.4, there is a new Enum data type, which you can use as:

class EventType(Enum):

    def __init__(self, id, code, message):
        self.id = id
        self.code = code
        self.message = message

    CANCELLED = 1, 'code1', 'message1'
    ONTIME = 2, 'code2', 'message2'
    DELAYED = 3, 'code3', 'message3'

To use, simply:

EventType.CANCELLED.code # returns 'message1'

If it's just the use case you described, you can also use named tuple if you want:

from collections import namedtuple

    Event = namedtuple('event', ['id', 'code', 'message'])
    
    class EventType:
        CANCELLED = Event(1, 'code1', 'message1')
        ONTIME = Event(2, 'code2', 'message2')
        DELAYED = Event(3, 'code3', 'message3')

To use this one:

EventType.CANCELLED # return event(id=1, code='code1', message='message1')
EventType.CANCELLED.message # return message1

Upvotes: 2

Ethan Furman
Ethan Furman

Reputation: 69021

If you want the value of the Enum member to be 0, 1, or 2, then you will need to either override __new__, or use aenum.

Using the aenum library1:

from aenum import MultiValueEnum

class EventType(MultiValueEnum):
    _init_ = 'value status message'
    CANCELLED = 0, 'status1', 'message1'
    ONTIME = 1, 'status2', 'message2'
    DELAYED = 2, 'status3', 'message3'

and in use:

>>> print(EventType.ONTIME)
EventType.ONTIME
>>> print(EventType.ONTIME.value)
1
>>> print(EventType.ONTIME.status)
status2
>>> print(EventType.ONTIME.message)
message2

See this answer for an example of how to do this with the stdlib Enum (which uses the overriding __new__ technique).

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

Upvotes: 5

Blckknght
Blckknght

Reputation: 104682

You can use a tuple as the value of your Enum objects, and include the status and message strings together with the integer code. Then you can add methods (or property descriptors) to fetch individual pieces of the tuple via the value attribute:

class EventType(Enum):
   CANCELLED = 0, "status1", "message1"
   ONTIME = 1, "status2", "message2"
   DELAYED = 2, "status3", "message3"

   @property
   def code(self):
       return self.value[0]

   @property
   def status(self):
       return self.value[1]

   @property
   def message(self):
       return self.value[2]

print(EventType.CANCELLED.message) # prints "message3"

Upvotes: 1

Related Questions