方法重载与类型注解

在 Python 中，方法重载是一个比较常见的需求，尤其当我们需要根据不同类型的参数执行不同的逻辑时。虽然 Python 本身并不支持传统意义上的方法重载，但我们可以通过类型注解、条件判断以及 @overload​ 装饰器来实现类似的功能。本文将通过一个简单的例子来演示如何处理不同参数类型的重载。

场景介绍

假设我们有一个操作处理器 ActionHandler​，该处理器根据不同的操作类型（如创建集群、扩展集群）来执行对应的逻辑。不同操作需要的参数并不相同，例如创建操作需要集群名称和节点列表，而扩展操作则需要集群 ID 和节点列表。

为了实现这一需求，我们可以通过 pydantic​ 定义数据模型，并使用类型注解与 @overload​ 来实现方法的重载。

代码解析

首先，我们定义了两个操作参数模型：CreateActionParams​ 和 ExpandActionParams​，分别对应“创建”和“扩展”操作。每个模型都通过 Literal​ 来指定 action​ 字段的固定值，以确保操作类型的唯一性和准确性。

from typing_extensions import Literal
from pydantic import BaseModel


class CreateActionParams(BaseModel):
    action: Literal["create"]
    name: str
    access_address: str
    node_ids: list[int]


class ExpandActionParams(BaseModel):
    action: Literal["expand"]
    cluster_id: int
    node_ids: list[int]

接下来，定义了一个通用的 Params​ 类，它包含了两种操作类型中的一种。这个类帮助我们将不同类型的参数模型统一起来，以便传递给后续的逻辑处理。

class Params(BaseModel):
    value: CreateActionParams | ExpandActionParams

方法重载

在 ActionHandler​ 类中，我们希望根据传入的参数类型，自动调用不同的处理方法。例如，当参数为 CreateActionParams​ 时，我们调用 _create​ 方法；当参数为 ExpandActionParams​ 时，则调用 _expand​ 方法。

为此，我们使用了 @overload​ 装饰器声明了不同参数类型的重载方法。实际的 run​ 方法使用 isinstance​ 来判断参数的类型，并调用对应的方法。

from typing import overload

class ActionHandler:

    @overload
    def run(self, params: CreateActionParams) -> None: ...

    @overload
    def run(self, params: ExpandActionParams) -> None: ...

    def run(self, params: CreateActionParams | ExpandActionParams) -> None:
        if isinstance(params, CreateActionParams):
            return self._create(params)
        elif isinstance(params, ExpandActionParams):
            return self._expand(params)
        else:
            raise Exception("Action not supported")

    def _create(self, params: CreateActionParams) -> None:
        print(f"{params.action} cluster {params.name}, use nodes: {params.node_ids}")

    def _expand(self, params: ExpandActionParams) -> None:
        print(f"{params.action} cluster {params.cluster_id}, use nodes: {params.node_ids}")

使用示例

我们可以通过如下方式使用 ActionHandler​ 来执行不同的操作。首先，我们定义了一个 action​ 函数，该函数接收一个 Params​ 对象，并将其传递给 ActionHandler​ 的 run​ 方法进行处理。

def action(params: Params) -> None:
    ActionHandler().run(params.value)

最后，我们通过 pydantic​ 的 model_validate​ 方法将字典转换为 Params​ 对象，并执行对应的逻辑：

if __name__ == "__main__":
    params = {
        "value": {
            "action": "create",
            "name": "test",
            "access_address": "127.0.0.1",
            "node_ids": [1, 2],
        }
    }
    action(Params.model_validate(params))

    params2 = {"value": {"action": "expand", "cluster_id": 1, "node_ids": [1, 2]}}
    action(Params.model_validate(params2))

在运行这段代码时，系统会根据 action​ 字段的值来判断应执行创建还是扩展操作。例如，action​ 为 "create"​ 时，输出为：

create cluster test, use nodes: [1, 2]

当 action​ 为 "expand"​ 时，输出为：

expand cluster 1, use nodes: [1, 2]

总结

通过使用 @overload​ 和类型注解，我们能够在 Python 中实现类似方法重载的功能，从而根据不同类型的参数执行不同的逻辑。本文展示了如何结合 pydantic​ 和 typing​ 模块来实现这一需求。虽然 Python 不支持传统的重载机制，但通过合理的设计，我们仍然可以实现灵活的、类型安全的代码。

这种方式特别适合需要处理多种不同类型的输入、并且希望在编译时捕获潜在错误的场景。

‍