如何为偏应用后的Python函数添加精确的类型注解？

需求场景 #

我们有类似FastAPI端点的函数，包含普通参数和抽象类类型的Service参数，希望复用该函数，通过注入不同的Service实现类生成多个版本，且注入顺序不影响结果。

示例函数：

def foobar(foo: int, bar: str, spam: SpamService) -> str:
    return spam.serve(foo, bar)

期望的复用方式：

foobar_rabbit = inject(foobar, RabbitService)
foobar_snake = inject(foobar, SnakeService)
foobar_rabbit(1, "rabot")
foobar_snake(2, "sniky")

对于包含多个依赖的函数：

def foobar(foo: int, egg: EggService, spam: SpamService) -> str:
    return spam.serve(foo, egg.do_stuff())

两种注入顺序应生成等价的函数：

foobar_1 = inject(foobar, SpamService1, EggService2)
foobar_1_ = inject(foobar, EggService2, SpamService1)  # 结果与foobar_1等价

已实现的无依赖代码 #

import abc
import functools
import inspect
import typing


class Service(abc.ABC):
    ...


class ServiceA(Service):
    @staticmethod
    @abc.abstractmethod
    def method_a(a: int) -> str:
        """
        该方法执行某些操作
        """


class ServiceA1(ServiceA):
    @staticmethod
    def method_a(a: int) -> str:
        return f"A1: {a}"


def inject(
        func: typing.Callable,
        *services: typing.Type[Service]
) -> functools.partial:
    annotations = inspect.get_annotations(func)
    del annotations["return"]

    bind_services = {
        key: service
        for key, value in annotations.items()
        if issubclass(value, Service)
        for service in services
        if issubclass(service, value)
    }

    return functools.partial(func, **bind_services)


def foobar(foo: int, spam: ServiceA) -> str:
    return spam.method_a(foo)


foobar_A1 = inject(foobar, ServiceA1)

if __name__ == '__main__':
    print(foobar_A1(1))  # 输出: A1: 1

问题描述 #

foobar_A1的类型签名不符合预期：调用时不传入参数，PyCharm不会给出警告，mypy也无法检测出错误。尝试过使用typing.TypeVar但无效，例如：

_SERVICE = typing.TypeVar("_SERVICE", bound=Service)
_RETURN = typing.TypeVar("_RETURN")

def inject(
        func: typing.Callable[[..., _SERVICE], _RETURN],
        *services: typing.Type[Service]
) -> functools.partial[typing.Callable[[_SERVICE, ...], _RETURN]]:

mypy会报错，且无法生成期望的(foo: int) -> str签名。

解决方案 #

要实现正确的类型注解，需利用Python的ParamSpec跟踪原函数参数，结合TypeVar区分Service类型与普通参数，修正后的实现如下：

import abc
import functools
import inspect
import typing
from typing import ParamSpec, TypeVar, Callable, Type

# 定义参数泛型和类型变量
P = ParamSpec("P")
R = TypeVar("R")
S = TypeVar("S", bound="Service")

class Service(abc.ABC):
    ...


class ServiceA(Service):
    @staticmethod
    @abc.abstractmethod
    def method_a(a: int) -> str:
        """
        该方法执行某些操作
        """


class ServiceA1(ServiceA):
    @staticmethod
    def method_a(a: int) -> str:
        return f"A1: {a}"


def inject(
    func: Callable[P, R],
    *services: Type[S]
) -> Callable[P, R]:
    annotations = inspect.get_annotations(func)
    del annotations["return"]

    bind_services = {
        key: service
        for key, value in annotations.items()
        if issubclass(value, Service)
        for service in services
        if issubclass(service, value)
    }

    # 生成partial并转换为正确的Callable类型
    partial_func = functools.partial(func, **bind_services)
    return typing.cast(Callable[P, R], partial_func)


def foobar(foo: int, spam: ServiceA) -> str:
    return spam.method_a(foo)


foobar_A1 = inject(foobar, ServiceA1)

if __name__ == '__main__':
    print(foobar_A1(1))  # 输出: A1: 1

关键说明 #

1. ParamSpec的作用：ParamSpec(P)用于精确捕获原函数的参数列表，TypeVar(R)捕获返回值类型，这样类型检查工具能清晰识别注入后剩余的参数。
2. 类型转换：functools.partial本身的类型不够精确，使用typing.cast将其转换为Callable[P, R]，让mypy和IDE能识别正确的函数签名。
3. Service匹配逻辑：保持原有的依赖匹配逻辑，通过注解识别Service类型参数，自动匹配传入的实现类，注入顺序不影响结果。

修改后，foobar_A1的签名会被正确识别为(foo: int) -> str，调用时缺少参数会触发IDE警告和mypy错误。

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