为 Enum 属性键入注释

我有一段代码:

import enum




class Color(enum.Enum):
RED = '1'
BLUE = '2'
GREEN = '3'




def get_color_return_something(some_color):
pass

如果我假设我将从 Color enum (例如: Color.RED)接收一个 enum 属性，那么如何正确地将类型注释添加到这个函数的 some_color变量中？

60605 次浏览
小开 
def get_color_return_something(some_color: Color):
pass
小开
最佳答案

Type hinting the Color class should work:

def get_color_return_something(some_color: Color):
print(some_color.value)
小开

The following will work with Pyton 3.9/PyCharm

from enum import Enum
from typing import Optional, Union




class Color(Enum):
RED: int = 1
GREEN: int = 2




def guess_color(x: Union[Color.RED, Color.GREEN]) -> Optional[ValueError]:
if x == Color.RED:
print("Gotcha!")
else:
return ValueError(f"It's not {Color.RED}")




guess_color(Color.RED)
小开

Another strange syntactic workaround is to specify Enum members as the type of the Enum class using the forward-referencing syntax of quoting (per PEP 484):

from enum import Enum




class ETest(Enum):
EXAMPLE: 'ETest' = "example"  <--- forward referenced type




def example() -> ETest:
return ETest.EXAMPLE




print(type(ETest.EXAMPLE.value))


<class 'str'>

In the image below it's evident that the warnings highlighted in PyCharm are no longer present.

type-hint with reference to Enum class

For reference, here is a screenshot of PyCharm's grievance with specifying the EXAMPLE member as a <str> type as makes sense:

PyCharm pylint warning Enum Typing

I am not a fan or this approach but it does get rid of the warning.

小开

You can try to use an option with type hint Literal.

From official PEP8 documentation we know that:

Literal it's type that can be used to indicate to type checkers that the corresponding variable or function parameter has a value equivalent to the provided literal (or one of several literals)

So in case if you need to use some specific values for a function argument it will be one of the best options. But this approach will not work fully as we expected, because of the type of the Enum values. Each value will have a type of Enum class. This means that for the code example below we will be able to put Color.GREEN as the function argument. So such a solution will be just information for developers, but not a mandatory rule for a function argument.

class Color(enum.Enum):
RED = '1'
BLUE = '2'
GREEN = '3'


print(type(Color.RED)  # will return <enum 'Color'>

Code example:

from enum import Enum
from typing import Literal




class Color(Enum):
RED = '1'
BLUE = '2'
GREEN = '3'


def some_function(some_color: Literal[Color.RED, Color.BLUE]) -> None:
pass

The second option it's fully correct solution provided by @ibarrond from the post above with just a class type hint.

some_color: Color

So here you can choose the option to work with depending on your needs.

From my point of view we can try to specify possible Enum values for developers, to be more clear in our requirements for a function.


提交答案