更改Python中的字典值类型和类型提示

Question

修改字典中存储的值的类型时处理输入的最合适方法是什么？

例如，如果我通过 mypy 运行以下代码，它会抛出错误“...没有属性“size”[attr-defined]”。

from typing import Dict

import numpy as np
import numpy.typing as npt
import torch
import torchvision.transforms as T

def to_tensor(my_dict: Dict[str, npt.NDArray[np.uint8]]) -> Dict[str, torch.FloatTensor]:
    for key, val in my_dict.items():
        my_dict[key] = T.functional.to_tensor(val)
    return my_dict

test_dict = {"foo": np.random.random((3,10,10)), "bar": np.random.random((3, 10, 10))}
test_dict = to_tensor(test_dict)
print(test_dict['foo'].size())

谢谢！

Answer 1

一旦定义

test_dict = {"foo": np.random.random((3,10,10))}

就 mypy 而言，您已经将

test_dict

的类型完全定义为

dict[str, npt.NDArray[np.uint8]]

。 您可以为

to_tensor

的结果使用不同的名称

test_dict_tensor = to_tensor(test_dict)

如果您不需要

FloatTensor

上的属性（尽管很明显您在这里需要），您也可以使用

test_dict

的类型明确表示值可以是任一类型

test_dict: dict[str, npt.NDArray[np.uint8] | torch.FloatTensor] = (
    {"foo": np.random.random((3,10,10))}
)

注意

X | Y

是 python >= 3.10 中

Union[X, Y]

的替代表示法（并且在 python <= 3.9 if you have

from __future__ import annotations

中启用）。