如何在 Python 中为特定输入类型强制执行特定输出类型？答案

【问题标题】：How can I enforce a specific output type for a specific input type in Python?如何在 Python 中为特定输入类型强制执行特定输出类型？
【发布时间】：2023-02-02 23:37:30
【问题描述】：

我一直在学习如何使用 Python 类型提示，并且有一个特定的用例让我很头疼。

假设我有以下 Pydantic 模型：

from pydantic import BaseModel


class Horse(BaseModel):
    speed: str
    race_wins: int

class HorseWithHat(Horse):
    hat_color: str

class Snake(BaseModel):
    length: str
    poisonous: bool

class SnakeWithHat(Snake):
    hat_color: str

# Etc.

我有各种其他动物模型，每个模型都有一个相关的带帽子的动物模型。我现在想实现一个给动物戴帽子的功能。类型签名类似于

def give_hat(animal: Animal, hat_color: str) -> AnimalWithHat

其中 Animal = Union[Horse, Snake, etc.] 和 AnimalWithHat = Union[HorseWithHat, SnakeWithHat, etc.]。当然，这个想法的问题是Horse可以进去，SnakeWithHat可以出来；我想强制一致性。

我的另一个想法是创建一个 WithHat 泛型。类型签名将是

def give_hat(animal: AnimalTypeVar, hat_color: str) -> WithHat[AnimalTypeVar]

AnimalTypeVar 是由Animal = Union[Horse, Snake, etc.] 绑定的类型变量。这将具有压缩重复的 WithHat 模型定义的优势，但是，我还没有想出一种方法来定义以这种方式工作的泛型（向输入类型添加单个属性）。

我希望我错过了一些简单的东西！有什么建议么？

（我知道我可以将非帽子和帽子模型结合起来，使 hat_color 成为可选属性，但在我的实际项目中，处理起来很挑剔。如果可能的话，我想要一个具有明显无帽和有帽的解决方案楷模。）

【问题讨论】：

我们在这里谈论多少种不同的动物（可能有帽子）？个位数？或者可能有数百个？在您定义该功能时它们是否已知？
我们现在只说几个。它们都有独特的属性，这就是为什么它们需要不同的模型。是的，在我定义函数时它们都已为人所知。

标签： python generics inheritance python-typing pydantic

【解决方案1】：

优雅的解决方案将当然是 intersection 类型，但我们在 Python 类型系统中还没有它们（目前）。这已经在 SO 上的一些帖子中进行了讨论，例如，请参阅here。

然后我们只需为WithHat定义一个protocol，为animal定义一个Animal类型变量（如您所述），并将返回类型定义为WithHat和Animal的交集。唉...

但是既然你提到只有几个动物模型并且你提前知道它们，你可以求助于用typing.overload重载give_hat的签名。

from typing import Union, overload
from typing_extensions import TypeAlias


Animal: TypeAlias = Union[Horse, Snake, Cow]
AnimalWithHat: TypeAlias = Union[HorseWithHat, SnakeWithHat, CowWithHat]


@overload
def give_hat(animal: Snake, hat_color: str) -> SnakeWithHat:
    ...


@overload
def give_hat(animal: Cow, hat_color: str) -> CowWithHat:
    ...


@overload
def give_hat(animal: Horse, hat_color: str) -> HorseWithHat:
    ...


def give_hat(animal: Animal, hat_color: str) -> AnimalWithHat:
    return NotImplemented  # actual implementation here

这不是很好，如果动物的数量开始增加，显然会在你的脸上爆炸，但它有效。

【讨论】：