mypy-arg应该是一组函数中的一个

Question

from enum import Enum
def strategy_a(x: float, y: float) -> float:
return x + y
def strategy_b(x: float, y: float) -> float:
return x * y
class Strategy(Enum):
A = strategy_a
B = strategy_b
def run_strategy(x: float, y: float, strategy: Strategy) -> float:
return strategy(x, y)

假设我有这样的东西，run_strategy的参数strategy接受一些可能的函数集。我怎么能键入它，这样只有那些函数才能传递，而mypy不会抛出错误。

注意：当mypy抱怨Strategy不可调用时，上面的代码抛出了一个错误。

然后将上述代码作为运行

run_strategy(5, 17, Strategy.A)

Answer 1

绕过将指定一个协议Strategy，该协议规定了调用的实现，并将Enum项重新命名为继承该协议，请参见示例blah.py:

from enum import Enum
from typing import Protocol

def strategy_a(x: float, y: float) -> float:
return x + y

def strategy_b(x: float, y: float) -> float:
return x * y

class Strategy(Protocol):
def __call__(self, x: float, y: float) -> float:
...

class StrategyChoice(Enum, Strategy):
A = strategy_a
B = strategy_b

def run_strategy(x: float, y: float, strategy: StrategyChoice) -> float:
return strategy(x, y)

mypy ./blah.py
Success: no issues found in 1 source file