Amin Merati
Amin Merati

Reputation: 2291

writing __str__ method for Enum class python

I have an enum class in python using list of strings:

from enum import Enum
myEnum = Enum('myEnum',['typeA',typeB','typeC'])

and I want to add method "str" to this enum class:

def __str__(self):
    return self.name

so that when I use str() function, I only get the 'typeX' part, i.e. the enum part withou class name. e.g.

print(str(myEnum.typeA)) # I want this to print "typeA" instead of "myEnum.typeA"

I do not know how to add the "str" method to this class as it is not class definition? Thanks for your help.

Upvotes: 2

Views: 7268

Answers (3)

bitinerant
bitinerant

Reputation: 1561

Here's another approach which I feel is simpler and easier to customize. This is based on the Allowed members and attributes of enumerations, part of the official Enum documentation.

from enum import Enum, auto

class myEnum(Enum):
    typeA = auto()
    typeB = auto()
    typeC = auto()

    def __str__(self):
        return self.name

Here is a usage example:

>>> print(myEnum.typeB)
typeB

Upvotes: 1

Ethan Furman
Ethan Furman

Reputation: 69041

If this is something you do a lot, you can make your own base Enum class, and then use that in the function calls:

class MyBaseEnum(Enum):
    def __str__(self):
        return self.name

myEnum = MyBaseEnum('myEnum',['typeA','typeB','typeC'])

and in use:

>>> print(myEnum.typeB)
typeB

Upvotes: 9

Amin Merati
Amin Merati

Reputation: 2291

Thanks for both the answers. I used it as below and it worked:

from enum import Enum
myEnum = Enum('myEnum',['typeA','typeB','typeC'])
def __str__(self):
    return self.name
myEnum.__str__ = __str__

st = str(myEnum.typeA)
print(st)

it printed "typeA" instead of "myEnum.typeA" which is what I needed!

Upvotes: 1

Related Questions