如何提升 defaultdict 的类型安全 - 给 key 增加类型

一般情况，我们只会给 defaultdict 的 value 增加类型，这样 defaultdict 就知道如何初始化。

但是有些情况，如果不给 key 增加类型约束容易写错，导致 bug。

给个案例：输入一行字符，分别统计出其中英文字母、空格、数字和其它字符的个数。

示例：'123runoobc kdf235*(dfl'

输出 char = 13,space = 2,digit = 6,others = 2

第一版

from collections import defaultdict
from typing import TypedDict


class Result(TypedDict):
    char: int
    space: int
    digit: int
    others: int




def count_category(string: str) -> Result:
    result = defaultdict(int)

    for char in string:
        if char.isalpha():
            result["chars"] += 1
        elif char.isspace():
            result["space"] += 1
        elif char.isdigit():
            result["digit"] += 1
        else:
            result["others"] += 1

    return Result(**result)



# char = 13,space = 2,digit = 6,others = 2
print(count_category("123runoobc  kdf235*(dfl"))

输出 {'digit': 6, 'chars': 13, 'space': 2, 'others': 2}

看到小问题了么，将 char 写成了 chars 但是IDE 并未提醒，即使我们使用了 TypedDict！

要是 defaultdict 支持指定 key 的范围就好了。

我们可以通过 DefaultDict[KeyType, ValueType] 来实现。

from collections import defaultdict
from typing import DefaultDict, Literal, TypeAlias, TypedDict


class Result(TypedDict):
    char: int
    space: int
    digit: int
    others: int


KeyType: TypeAlias = Literal["char", "space", "digit", "others"]


def count_category(string: str) -> Result:
    result: DefaultDict[KeyType, int] = defaultdict(int)

    for char in string:
        if char.isalpha():
            result["char"] += 1
        elif char.isspace():
            result["space"] += 1
        elif char.isdigit():
            result["digit"] += 1
        else:
            result["others"] += 1

    return Result(**result)

# 预期 char = 13,space = 2,digit = 6,others = 2
print(count_category("123runoobc  kdf235*(dfl"))
# {'digit': 6, 'char': 13, 'space': 2, 'others': 2}

当我们尝试 result["chars"] += 1 会报错提醒 Type "Literal['charx']" is not assignable to type "KeyType" 🎉🎉🎉。