如何通过Python装饰器添加输入参数类型并实现正确类型提示

一、修复当前示例的类型提示问题 #

你的代码中类型提示报错，是因为装饰后的函数类型注解仍保留原函数的int | float，类型检查器不知道它现在支持str类型。以下两种方法可以解决：

方法1：用@overload声明多类型签名 #

直接给装饰后的函数添加重载签名，明确标注支持的所有参数类型：

import functools
from typing import overload

def str_is_ok_too(func):
    '''装饰器：处理字符串类型的输入'''
    @functools.wraps(func)
    def wrapper(arg):
        if isinstance(arg, str):
            print("input is a string")
            print(f'"{arg}"')
        else:
            print("input is int or float")
            func(arg)
    return wrapper

@overload
def accept_int_or_float(arg: int) -> None: ...
@overload
def accept_int_or_float(arg: float) -> None: ...
@overload
def accept_int_or_float(arg: str) -> None: ...

@str_is_ok_too
def accept_int_or_float(arg: int | float) -> None:
    '''打印int或float类型的输入'''
    print(arg)

accept_int_or_float(1)       # 正常
accept_int_or_float(1.3)     # 正常
accept_int_or_float("5")     # 类型提示不再报错

方法2：给装饰器添加类型标注并更新注解 #

通过修改装饰器的类型定义，同时手动更新函数的参数注解，让类型检查器自动识别扩展类型：

import functools
from typing import ParamSpec, TypeVar, Callable

P = ParamSpec('P')
R = TypeVar('R')

def str_is_ok_too(func: Callable[P, R]) -> Callable[P | tuple[str], R | None]:
    '''装饰器：处理字符串类型的输入'''
    @functools.wraps(func, updated=("__annotations__",))
    def wrapper(arg):
        if isinstance(arg, str):
            print("input is a string")
            print(f'"{arg}"')
            return None
        else:
            print("input is int or float")
            return func(arg)
    # 更新参数注解，添加str类型
    wrapper.__annotations__['arg'] = int | float | str
    return wrapper

@str_is_ok_too
def accept_int_or_float(arg: int | float) -> None:
    '''打印int或float类型的输入'''
    print(arg)

accept_int_or_float(1)       # 正常
accept_int_or_float(1.3)     # 正常
accept_int_or_float("5")     # 类型提示正常

二、基于ABC实现自定义类型的"重载"支持 #

如果要支持自定义抽象基类（ABC）的类型分派，推荐使用functools.singledispatch——它原生支持根据参数类型匹配处理函数，且能自动识别ABC的子类，完美实现类似重载的效果，同时类型提示完全正常。

示例代码 #

from abc import ABC, abstractmethod
from functools import singledispatch

# 定义抽象基类
class CustomType(ABC):
    @abstractmethod
    def process(self) -> str:
        pass

# 实现两个子类
class TypeA(CustomType):
    def process(self) -> str:
        return "Processing TypeA"

class TypeB(CustomType):
    def process(self) -> str:
        return "Processing TypeB"

# 定义单分派泛函数
@singledispatch
def handle_input(arg):
    raise NotImplementedError(f"Unsupported type: {type(arg)}")

# 注册基础类型的处理逻辑
@handle_input.register
def _(arg: int | float) -> None:
    print(f"Handling number: {arg}")

@handle_input.register
def _(arg: str) -> None:
    print(f"Handling string: \"{arg}\"")

# 注册抽象基类，所有子类自动匹配此逻辑
@handle_input.register(CustomType)
def _(arg: CustomType) -> None:
    print(arg.process())

# 测试调用
handle_input(10)          # 输出: Handling number: 10
handle_input("hello")     # 输出: Handling string: "hello"
handle_input(TypeA())     # 输出: Processing TypeA
handle_input(TypeB())     # 输出: Processing TypeB

优势 #

• 无需手动写isinstance判断，逻辑更简洁
• 原生支持ABC，子类自动继承匹配规则
• 类型提示完全兼容，类型检查器能识别每种参数对应的处理逻辑
• 扩展性强，新增类型只需注册对应的处理函数

三、概念测试的建议 #

1. 优先用singledispatch实现类型分派，比自定义装饰器更易维护，类型提示友好
2. 若必须自定义装饰器，一定要更新函数的__annotations__或补充@overload签名，确保类型检查器识别扩展类型
3. 对于ABC支持，singledispatch原生支持注册抽象基类，无需额外处理
4. 如需多参数组合的重载，可结合typing.overload或第三方库multipledispatch实现

内容的提问来源于stack exchange，提问作者gero