how to make an enum point to a class

user_12475905 0
Viewed 144

Suppose i have this enum for storing all the different Entity types

class Entities(Enum):
    TREE = auto()
    ROCK = auto()
    FLOWER = auto()

I would like to create a function that takes one of these (TREE, ROCK...) enums, and knows that an enum corresponds to a class i have. eg.:

def myFunc(EntityType):
    return type(EntityType)

print(myFunc(Entities.ROCK))
>>>ROCK (where ROCK is an instance of the ROCK class)

and if there is a way to do that, is there one, to maybe even initialise the class eg.:

def myFunc(EntityType):
    myObj = EntityType(pos=(0,0))
    return myObj
enums class python
3 Answers

What if you just ditch the auto and use the classes themselves as the values of Entities? Suppose Tree, Rock, and Flower are the names of your classes:

class Entities(Enum):
    TREE = Tree
    ROCK = Rock
    FLOWER = Flower

Here Entities.TREE.value is the class constructor for Tree.

user_108633270

Here an example to Kyle Parsons' answer:

from enum import Enum
from dataclasses import dataclass

@dataclass
class Animal:
    name: str
    age: int
    type: str = None

@dataclass
class Cat(Animal):
    type: str = 'Cat'

@dataclass
class Dog(Animal):
    type: str = 'Dog'


class AnimalType(Enum):
    DOG = Dog
    CAT = Cat


def get_animal(type: Enum, name: str, age: int):
    return type.value(name, age)

print(get_animal(AnimalType.CAT, 'Peter', 12))
user_135563030

You can add attributes to an Enum object, or you can map an Enum using a dict. There are other options as well, but these seem simplest.

Let's say you have classes Tree, Rock, Flower, etc, corresponding to the values of the Enum:

class Tree:
    def __init__(self, height, coords):
        pass

class Rock:
    def __init__(self, coords):
        pass

class Flower:
    def __init__(self, color, iscarnivore, coords):
        pass

I've specifically shown an extended version, where each class has a different initializer, and a different set of defaults. If they are all the same, use the existing answers.

Option 1 is to define the enum like this:

class Entities(Enum):
    TREE = (Tree, 100, (0, 0))
    ROCK = (Rock, (0, 0))
    FLOWER = (Flower, 'red', True, (0, 0))

    def __new__(cls, t, *args):
        obj = object.__new__(cls)
        obj._value_ = len(cls.__members__) + 1
        obj.type = t
        obj.defaults = args
        return obj

    def init(self):
        return self.type(*self.defaults)

Now, my_func is just the init method of the enum itself:

>>> FLOWER.init() # Calls Flower('red', False, (0, 0))

The second option would be to map the Enum members to the class:

cmap = {
    Entitites.TREE: (Tree, 100, (0, 0)),
    Entitites.ROCK: (Rock, (0, 0)),
    Entitites.FLOWER: (Flower, 'red', True, (0, 0)),
}

def my_func(entity):
    t, *args = cmap[entity]
    return t(*args)
user_29887300
Related