Python泛型约束：依赖类型组合实现

本文介绍了如何在Python中使用泛型和协议(Protocol)来实现更精确的类型提示，特别是当泛型类型之间存在依赖关系时。通过定义一个Indexable协议，并结合TypeVar和Generic，可以约束ApplyTo类，使其能够根据to参数的类型，正确地推断出data参数的类型，从而提高代码的类型安全性。

在Python中，类型提示可以显著提高代码的可读性和可维护性。然而，在处理具有复杂类型依赖关系的场景时，简单的类型提示可能不足以表达代码的意图。本文将介绍如何使用泛型和协议(Protocol)来解决这类问题，并以一个实际的例子进行说明。

使用 Protocol 定义可索引类型

当需要表达一个类型必须支持某种操作，但又不关心其具体的实现时，可以使用Protocol。 在本例中，我们需要一个类型，它支持通过键进行索引和赋值操作。 我们可以定义一个名为Indexable的协议如下：

import typing

K = typing.TypeVar('K', contravariant=True)
class Indexable(typing.Protocol[K]):
    def __getitem__(self, key: K):
        pass

    def __setitem__(self, key: K, value: typing.Any):
        pass

这里，Indexable协议定义了__getitem__和__setitem__方法，这意味着任何实现了这两个方法的类都符合Indexable协议。K是一个类型变量，表示键的类型，contravariant=True 表示 K 是逆变的，允许子类型替换父类型。

使用泛型约束参数类型

现在，我们可以使用Indexable协议来约束ApplyTo类的data参数的类型。 首先，定义一个类型变量H，它表示to参数的类型，并且必须是可哈希的：

H = typing.TypeVar('H', bound=typing.Hashable)
DispatchType = typing.Literal['separate', 'joint']
class ApplyTo(typing.Generic[H]):
    _to: typing.Sequence[H]
    _dispatch: DispatchType
    _transform: typing.Callable[..., typing.Any]  # TODO Initialize `_transform`

    def __init__(self, to: typing.Sequence[H] | H, dispatch: DispatchType = 'separate') -> None:
        self._dispatch = dispatch
        self._to = to if isinstance(to, typing.Sequence) else [to]

    def __call__(self, data: Indexable[H]) -> typing.Any:
        if self._dispatch == 'separate':
            for key in self._to:
                data[key] = self._transform(data[key])
            return data

        if self._dispatch == 'joint':
            args = [data[key] for key in self._to]
            return self._transform(*args)

        assert False

在这里，ApplyTo类被定义为泛型类，它接受一个类型参数H，H被限制为Hashable类型。 __init__方法接受一个to参数，它的类型是H或H的序列。 __call__方法接受一个data参数，它的类型是Indexable[H]，这意味着data必须是一个可索引的类型，并且它的键的类型必须是H。

使用示例

def main() -> None:
    r0 = ApplyTo(to=0)([1, 2, 3])
    # typechecks
    r0 = ApplyTo(to=0)({1: 'a', 2: 'b', 3: 'c'})
    # typechecks

    r1 = ApplyTo(to='a')(['b', 'c', 'd'])
    # does not typecheck: Argument 1 to "__call__" of "Applier" has incompatible type "list[str]"; expected "Indexable[str]"
    r1 = ApplyTo(to='a')({'b': 1, 'c': 2, 'd': 3}) 
    # typechecks

在这个例子中，ApplyTo(to=0)可以接受一个列表或一个字典，因为整数既可以作为列表的索引，也可以作为字典的键。 然而，ApplyTo(to='a')只能接受一个字典，因为字符串不能作为列表的索引。

总结

通过使用泛型和协议(Protocol)，我们可以更精确地约束Python代码中的类型，从而提高代码的类型安全性。 在处理具有复杂类型依赖关系的场景时，这种方法尤其有用。

注意事项:

• Protocol 是 Python 3.8 引入的，如果使用更早的版本，需要安装 typing_extensions 库。
• 类型提示只能在运行时提供类型信息，并不能阻止运行时错误。 类型提示的主要作用是帮助静态类型检查器（如mypy）在代码运行之前发现潜在的类型错误。