1

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!

Alex Waygood
  • 6,304
  • 3
  • 24
  • 46
Andrew Stewart
  • 173
  • 6

1 Answers1

1

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).

Alex Waygood
  • 6,304
  • 3
  • 24
  • 46
joel
  • 6,359
  • 2
  • 30
  • 55
  • Thanks Joel. Is it "unpythonic" to modify the types in place? I understand mypy's strict typing but should I add an ignore or should I refactor the code? I assume there is some memory impact if I instantiate two dictionaries, ie. test_dict and test_dict_tensor? – Andrew Stewart Sep 13 '21 at 16:03
  • Also, yes, I can't wait to transition to 3.9+ to not need so many typing imports. Thanks! – Andrew Stewart Sep 13 '21 at 16:05
  • 1
    @AndrewStewart I don't know if it's unpythonic. Since you're using mypy it would seem a waste and perhaps a little confusing for future readers to `type: ignore` it. I usually only resort to that if mypy can't express what I mean. I can't comment on the memory impact. – joel Sep 13 '21 at 17:16