CODEWITHSUNDEEP
pythonpython-typingmypy
Andrew Stewart
Andrew Stewart

Reputation: 193

Changing dictionary value types in python and Type hinting

What is the most appropriate way to handle typing when modifying the type of values stored in a dictionary?

For example, if I run the following code through mypy, it throws an error '... has no attribute "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())

Thanks!

Upvotes: 1

Views: 617

Answers (1)

joel
joel

Reputation: 7867

Once you define

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

you've completely defined the type of test_dict as dict[str, npt.NDArray[np.uint8]] as far as mypy is concerned. You can use a different name for the result of to_tensor

test_dict_tensor = to_tensor(test_dict)

If you didn't need attributes on FloatTensor (though it's clear you do here), you could alternatively be explicit with the type of test_dict that the values can be either type

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

Note X | Y is alternative notation for Union[X, Y] in python >= 3.10 (and in python <= 3.9 if you have from __future__ import annotations enabled).

Upvotes: 1

Related Questions