如何为函数装饰器的可选kwargs参数foo添加存在时的str类型约束？

需求说明 #

我实现了一个函数装饰器，它会检查可选关键字参数foo：如果foo存在，必须是str类型；其余的*args和**kwargs直接传递给被装饰函数。现在需要通过类型提示让静态检查工具（如mypy）识别出foo类型错误的情况。

当前问题 #

原代码使用ParamSpec定义参数类型，但mypy无法识别foo的类型错误，执行检查后无报错。

原代码 #

from typing import Callable, ParamSpec, TypeVar, TypedDict, Required, NotRequired, Unpack, cast, Any

FnRet = TypeVar("FnRet")

# how can this be partially constrained to _required_ type `str`,
# if `foo` item is present?
FnParams = ParamSpec("FnParams")

FnType = Callable[FnParams, FnRet]

def decorator(fn: FnType) -> FnType:
    def wrapper(*args: FnParams.args, **kwargs: FnParams.kwargs) -> FnRet: # type: ignore[type-var]
        if "foo" in kwargs:
            assert type(kwargs["foo"]) == str
        return cast(FnRet, fn(*args, **kwargs))
    return wrapper

## Should pass:

@decorator
def ok_1() -> int:
    return 42


class RequiredFooKwArgs(TypedDict):
    foo: Required[str]

@decorator
def ok_required_foo(*_args: Any, **kwargs: Unpack[RequiredFooKwArgs]) -> int:
    return 42

class NotRequiredFooKwArgs(TypedDict):
    foo: NotRequired[str]

@decorator
def ok_not_required_foo(*_args: Any, **kwargs: Unpack[NotRequiredFooKwArgs]) -> int:
    return 42


## The Badies:

class BadTypeRequiredFooKwArgs(TypedDict):
    foo: Required[int] # <--- this should be flagged as the wrong type

@decorator
def bad_type_required_foo(*_args: Any, **kwargs: Unpack[BadTypeRequiredFooKwArgs]) -> int:
    return 42


class BadTypeNotRequiredFooKwArgs(TypedDict):
    foo: NotRequired[int] # <--- this should be flagged as the wrong type

@decorator
def bad_type_not_required_foo(*_args: Any, **kwargs: Unpack[BadTypeNotRequiredFooKwArgs]) -> int:
    return 42

当前检查结果 #

执行命令：

mypy --enable-incomplete-feature=Unpack ~/tmp.py

输出：Success: no issues found in 1 source file

解决方案 #

ParamSpec无法直接对单个关键字参数添加条件约束，改用Protocol结合TypedDict来定义被装饰函数的类型，明确要求：如果kwargs中包含foo，则其类型必须为str，同时允许其他任意关键字参数。

修改后的代码 #

from typing import Callable, TypeVar, TypedDict, Required, NotRequired, Unpack, cast, Any, Protocol

FnRet = TypeVar("FnRet")

# 定义核心约束：foo可选，但如果存在必须是str
class FooAllowedKwargs(TypedDict):
    foo: NotRequired[str]

# 允许其他任意关键字参数，同时满足FooAllowedKwargs的约束
class AllowedKwargs(FooAllowedKwargs, total=False):
    [key: str]: Any

# 用Protocol定义被装饰函数的符合类型
class DecoratedFn(Protocol[FnRet]):
    def __call__(self, *args: Any, **kwargs: Unpack[AllowedKwargs]) -> FnRet:
        ...

def decorator(fn: DecoratedFn[FnRet]) -> DecoratedFn[FnRet]:
    def wrapper(*args: Any, **kwargs: Any) -> FnRet:
        if "foo" in kwargs:
            # 用isinstance代替type检查，更符合Python类型判断规范
            assert isinstance(kwargs["foo"], str)
        return cast(FnRet, fn(*args, **kwargs))
    return wrapper

## Should pass:

@decorator
def ok_1() -> int:
    return 42

class RequiredFooKwArgs(TypedDict):
    foo: Required[str]

@decorator
def ok_required_foo(*_args: Any, **kwargs: Unpack[RequiredFooKwArgs]) -> int:
    return 42

class NotRequiredFooKwArgs(TypedDict):
    foo: NotRequired[str]

@decorator
def ok_not_required_foo(*_args: Any, **kwargs: Unpack[NotRequiredFooKwArgs]) -> int:
    return 42

## The Badies:

class BadTypeRequiredFooKwArgs(TypedDict):
    foo: Required[int] # <--- 现在会被mypy标记错误

@decorator
def bad_type_required_foo(*_args: Any, **kwargs: Unpack[BadTypeRequiredFooKwArgs]) -> int:
    return 42

class BadTypeNotRequiredFooKwArgs(TypedDict):
    foo: NotRequired[int] # <--- 现在会被mypy标记错误

@decorator
def bad_type_not_required_foo(*_args: Any, **kwargs: Unpack[BadTypeNotRequiredFooKwArgs]) -> int:
    return 42

修改后检查结果 #

再次执行相同的mypy命令，会得到如下错误提示：

tmp.py:43: error: Argument 1 to "decorator" has incompatible type "Callable[[VarArg(Any), KwArg(BadTypeRequiredFooKwArgs)], int]"; expected "DecoratedFn[int]"  [arg-type]
tmp.py:51: error: Argument 1 to "decorator" has incompatible type "Callable[[VarArg(Any), KwArg(BadTypeNotRequiredFooKwArgs)], int]"; expected "DecoratedFn[int]"  [arg-type]
Found 2 errors in 1 source file (checked 1 source file)

说明 #

• FooAllowedKwargs定义了foo的核心类型约束：可选，但存在时必须是str。
• AllowedKwargs继承自FooAllowedKwargs并设置total=False，允许添加其他任意关键字参数，同时保留foo的约束。
• DecoratedFn协议确保被装饰函数的kwargs必须符合AllowedKwargs的规则，这样mypy就能静态检查出foo类型错误的情况。

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