Python type hint for class member with None type or dict

Question

I have class member which is either None or a dict. I have a method for the lengths of that member that returns None if the member is None or the size of the dictionary if it exists:

class A:
    def __init__(self, a: bool = False) -> None:
        self.has_a = a
        self.prop_a = {"a" : 1} if a else None

    def length_prop(self) -> int:
        length_prop = None
        if self.has_a:
            length_prop = len(self.prop_a)
        return length_prop

If I check the typing with mypy:

test.py:10: error: Argument 1 to "len" has incompatible type "Optional[Dict[str, int]]"; expected "Sized"
                length_prop = len(self.prop_a)
                                  ^
test.py:11: error: Incompatible return value type (got "Optional[int]", expected "int")
            return length_prop

Very understandable, since it does not know if self.prop_a is really dict. How can I design my code to get around this and get the right typing?

Answer 1

There are two issues:

1. mypy can't and shouldn't infer the relationship between has_a and prop_a across the functions. You need to include an explicit check for None inside the length_prop function. This error is not spurious -- for example, what happens if some external code modifies the value of has_a but not prop_a?

   The simplest solution is to do away with the has_a field, since all the information it encodes can be retrieved using a check on prop_a (for example, prop_a is not None) instead.

2. The return value of length_prop can be None, so use Optional[int] rather than int for the length_prop method.

The following code snippet resolves both of these issues. mypy does not find any errors in this class definition:

from typing import Optional

class A:
    def __init__(self, a: bool = False) -> None:
        self.prop_a = {"a" : 1} if a else None

    def length_prop(self) -> Optional[int]:
        length_prop = None
        if self.prop_a is not None:
            length_prop = len(self.prop_a)
        return length_prop