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 is not a callable。

以上代码则运行为

run_strategy(5, 17, Strategy.A)

Answer 1

一个变通方法是指定一个协议 Strategy 来规定 call 的实现，并将 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

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

mypy arg should be one of a set of functions

python

mypy

python-typing