Diogo
Diogo

Reputation: 733

Type annotation for mutable dictionary

I have a class MyClass with 2 methods. The methods change an attribute dic of type dict[str, A] by adding objects of subtypes B and C:

class A:
    ...

class B(A):
    def method_B(self): ...

class C(A):
    def method_C(self): ...

class MyClass:
    dic: dict[str, A] = {}
    
    def add_B_obj(self, name: str):
        self.dic[name] = B()

    def add_C_obj(self, name: str):
        self.dic[name] = C()

I would like to access the methods of the objects inside the dictionary after their creation, for instance, using Pylance.

For that, the following snippet doesn't work (see the warning):

m = MyClass()
m.add_B_obj('B_obj')

m.dic['B_obj'].method_B() # Warning: Cannot access member "method_B" for type "A"

To solve that, I tried to use typing.cast.
I'm not interested in objects of type C, so I used cast to force m.dic to be of type dict[str,B].
That allows me to access the method, but there was a new warning:

from typing import cast

m = MyClass()
m.add_B_obj('B_obj')

m.dic = cast(dict[str,B],m.dic['B_obj']) # WARNING

m.dic['B_obj'].method_B # now pylance shows me the method, but there is a new warning

The new warning is:

 Cannot assign member "dic" for type "MyClass"
  "dict[str, B]" is incompatible with "dict[str, A]"
    TypeVar "_VT@dict" is invariant
      "B" is incompatible with "A"

How do I solve this new warning?

Upvotes: 2

Views: 1345

Answers (1)

Daniil Fajnberg
Daniil Fajnberg

Reputation: 18683

Type safety for dict values

All values in a dictionary are treated the same by a static type checker.

If you tell the type checker that a dictionary dic can contain values of both type A and type B, but only B has the method method_B, then it logically follows that not all of the values in dic have method_B.

There is no way for the type checker to know for certain, if some particular value from dic has method_B.

In your example, B is a subtype of A, so by annotating dic with dict[str, A] you allow both A and B to be values in that dictionary. You could make the error more explicit, if you annotate it with dict[str, Union[A, B]] (which is equivalent from a typing perspective); then mypy for example will correctly tell you:

error: Item "A" of "Union[A, B]" has no attribute "method_B"

Highlighting the fact that B has it, but that is not enough for type safety.

The only way to assure the type checker that a specific value in the dictionary has method_B is to explicitly assert it for that specific key:

...
assert isinstance(m.dic['B_obj'], B)
m.dic['B_obj'].method_B()

Even this doesn't satisfy every type checker. mypy is happy, but the PyCharm built-in type checker still marks that last call saying:

Unresolved attribute reference 'method_B' for class 'A'

The other option is to use cast properly with an intermediary variable:

b = cast(B, m.dic['B_obj'])
b.method_B()

This should satisfy all type checkers.

I hope my explanation makes it clear that there is no other way around this.


Type (in)variance

As for the variance issue after cast-ing, your attempt is wrong for two reasons.

Firstly, you are only passing the value corresponding to the key 'B_obj' to cast, not the entire dic. Which means that after your assignment

m.dic = cast(dict[str, B], m.dic['B_obj'])

The attribute dic of object m is just that value, i.e. an instance of class B. This would lead to an TypeError, as soon as you tried this:

m.dic['B_obj']

Because B is not subscriptable (i.e. can not be used with []).

The type warning you then saw was because the type checker took you by your word and assumed you were really assigning something of type dict[str, B] to m.dic (even though you weren't).

And this is not allowed, because (as you are correctly told by the type checker) a dictionary is invariant in its value type. This is true by default for all generic types in Python, as stated in PEP 484.

In other words, since you declared .dic to be of type dict[str, A], you are not allowed to assign dict[str, B] to it, even though B is a subclass of A.


Niche alternative: TypedDict

If you know ahead of time that only very specific keys will be allowed in your dictionary, you can use the TypedDict construct.

It doesn't seem like this is what you are going for, since you allow adding keys via a public interface on MyClass, but if it were, you would just need to define your dictionary types as well as a corresponding Union of Literal types to annotate the name parameter in your methods. This amounts to a pretty verbose setup, but it works:

from typing import Any, Literal, TypeAlias, TypedDict, Union


class A:
    ...


class B(A):
    def method_B(self) -> Any: ...


class C(A):
    def method_C(self) -> Any: ...


class DicT(TypedDict, total=False):
    A_obj: A
    B_obj: B
    C_obj: C


KeyT: TypeAlias = Literal["A_obj", "B_obj", "C_obj"]


class MyClass:
    dic: DicT = {}

    def add_B_obj(self, name: KeyT) -> None:
        self.dic[name] = B()

    def add_C_obj(self, name: KeyT) -> None:
        self.dic[name] = C()


if __name__ == '__main__':
    m = MyClass()
    m.add_B_obj('B_obj')
    m.dic['B_obj'].method_B()

This passes mypy --strict with no issues.

Note that total=False on your TypedDict is necessary, since you want to allow any number of the supported keys to be present.

I also used TypeAlias just for better readability here.


In conclusion, I would recommend going the assert route as mentioned in the beginning. This not only satisfies some type checkers (don't know about Pylance), but it also adds an additional sanity check for you. Additional assertions are rarely a bad idea in my opinion.

Hope this helps.

Upvotes: 1

Related Questions