CODEWITHSUNDEEP
pythonenums
GhostCat
GhostCat

Reputation: 140633

How to write "correct" member function for enums?

See this code:

from __future__ import print_function
from enum import Enum

class AType(Enum):
    A = 1
    B = 2

    def as_key(self):
        if self.name == AType.A:
            return 'key for As'
        else:
            return 'some other key'

print(AType.A.as_key())
print(AType.B.as_key())

Coming from the java world; I would have expected it to print 

key for As
some other key

but I get:

some other key
some other key

Probably super-simply, but what is the correct way to write such a member function that gives me a result depending on the "enum constant" I invoke the method on?

Upvotes: 0

Views: 84

Answers (1)

Idan Arye
Idan Arye

Reputation: 12633

First, to get it out of the way: self.name is already key, so

def as_key(self):
    return self.name

But I assume this is not what you wanted - that function was probably just a basic example, and you want to write more complex functions that do different things based on the enum.

Enum values are singletonish - AType.A will always return the exact same object(as opposed to the string "A", which may create a new object each time - though as an optimization the literal will always return the same object if it's still in memory).

Also, inside the method, self is the actual enum object - AType.A or AType.B.

So - just compare self to the enum values:

def as_key(self):
    if self is AType.A:
        return 'A'
    elif self is AType.B:
        return 'B'
    else:
        assert False, 'Unknown AType %s' % (self,)

BTW: the is operator in Python is equivalent to == in Java - it checks for exact object identity(that it's the same object in the same place in memory). Python's == operator is equivalent to Java's equals() method - it can be customized to compare the values of the objects.

Upvotes: 3

Related Questions