How to anotate function Enums

Viewed 107

I have a function that takes other functions as argument, There are some possible options that the function can take, and I want to annotate them properly using Enums and type hints:

This is what I tried:

from typing import Callable
from enum import Enum


def test1():
    print('hi')

def test2():
    print('hi2')

class FuncEnum(Enum):
    FUN1 = test1
    FUN2 = test2

# option 1, mypy will complain because FuncEnum.FUN1 is not in the Callable type
def my_annotated_function(func: Callable = FuncEnum.FUN1):
    func()

# option 2, func is not callable. as it is FuncEnum type (not Callable)
def my_annotated_function(func: FuncEnum = FuncEnum.FUN1):
    func()

The code works the problem is only that the IDE/mypy raises a warning with both syntax.

I also tried sub-classing Callable, but it doesn't seem to work:

class FuncEnum(Callable, Enum):
    FUN1 = test1
    FUN2 = test2

Somehow the lintern/mypy seems to be lost with this kind of annotations, is there any work around?

Notice that the calls:

my_annotated_function(FuncEnum.FUN1) # should be a valid input

def random_function():
    pass

my_annotated_function(random_function) # should be raised by mypy
2 Answers

Enums don't work like that. If you do

class ExampleEnum(Enum):
    x = 1
    y = 2

ExampleEnum.x isn't 1. The enum metaclass backend sets ExampleEnum.x to an instance of ExampleEnum with a value attribute set to 1. The enum instance isn't its value, and it doesn't behave like its value.

Aside from that, there's an additional problem with trying to use functions as enum values. In your case,

class FuncEnum(Enum):
    FUN1 = test1
    FUN2 = test2

the enum backend thinks FUN1 and FUN2 are supposed to be methods rather than enumeration members. The enum backend doesn't support using functions as enum values. mypy will think FuncEnum.FUN1 is an instance of FuncEnum, but the actual runtime machinery won't handle things that way, so you'll get conflicts between runtime behavior and what static analysis thinks.

The annotation system doesn't support what you're trying to do, with or without enums. The closest thing is typing.Literal, but you can't use that with functions. You can use it with enums, but enums don't do what you're trying to do.

About using functions as enum members, see this question. If you want to allow only passing enum members, then the following will work (playground):

from functools import partial
from typing import Callable
from enum import Enum


def test1():
    print('hi')

def test2():
    print('hi2')

class FuncEnum(Enum):
    FUN1 = partial(test1)
    FUN2 = partial(test2)

def my_annotated_function2(func: FuncEnum = FuncEnum.FUN1):
    func.value()  # Actual callable is in `value` attribute
Related