How to handle inheritance in mypy?

Question

I'm trying to add mypy to my Python project but I have found a roadblock. Let's say I have the following inheritance:

class BaseClass:
    base_attribute: str


class A(BaseClass):
    attribute_for_class_A: str


class B(BaseClass):
    attribute_for_class_B: str

Now let's create some code that handle both instances of these classes, but without really knowing it:

@dataclass
class ClassUsingTheOthers:
    fields: Dict[str, BaseClass]

    def get_field(self, field_name: str) -> BaseClass:
        field = self.fields.get(field_name)
        if not field:
            raise ValueError('Not found')
        return field

The important bit here is the get_field method. Now let's create a function to use the get_field method, but that function will require to use a particlar subclass of BaseClass, B, for instance:

def function_that_needs_an_instance_of_b(instance: B):
    print(instance.attribute_for_class_B)

Now if we use all the code together, we can get the following:

if __name__ == "__main__":
    class_using_the_others = ClassUsingTheOthers(
        fields={
            'name_1': A(),
            'name_2': B()
        }
    )
    function_that_needs_an_instance_of_b(class_using_the_others.get_field('name_2'))

Obviously, when I run mypy to this file (in this gist you find all the code), I get the following error, as expected:

 error: Argument 1 to "function_that_needs_an_instance_of_b" has incompatible type "BaseClass"; expected "B"  [arg-type]

So my question is, how do I fix my code to make this error go away? I cannot change the type hint of the fields attribute because I really need to set it that way. Any ideas? Am I missing something? Should I check the type of the field returned?

Answer 1

I cannot change the type hint of the fields attribute

Well, there is your answer. If you declare fields to be a dictionary with the values type BaseClass, how do you expect any static type checker to know more about it?

(Related: Type annotation for mutable dictionary)

The type checker does not distinguish between different values of the dictionary based on any key you provide.

If you knew ahead of time, what the exact key-value-pairs can be, you could either do this with a TypedDict (as @dROOOze suggested) or you could write some ugly overloads with different Literal string values for field_name of your get_field method.

But none of those apply due to your restriction.

So you are left with either type-narrowing with a runtime assertion (as alluded to by @juanpa.arrivillaga), which I would recommend, or placing a specific type: ignore[arg-type] comment (as mentioned by @luk2302) and be done with it.

The former would look like this:

from dataclasses import dataclass

class BaseClass:
    base_attribute: str

@dataclass
class A(BaseClass):
    attribute_for_class_A: str

@dataclass
class B(BaseClass):
    attribute_for_class_B: str

@dataclass
class ClassUsingTheOthers:
    fields: dict[str, BaseClass]

    def get_field(self, field_name: str) -> BaseClass:
        field = self.fields.get(field_name)
        if not field:
            raise ValueError('Not found')
        return field

def function_that_needs_an_instance_of_b(instance: B) -> None:
    print(instance.attribute_for_class_B)

if __name__ == '__main__':
    class_using_the_others = ClassUsingTheOthers(
        fields={
            'name_1': A(attribute_for_class_A='foo'),
            'name_2': B(attribute_for_class_B='bar'),
        }
    )
    obj = class_using_the_others.get_field('name_2')
    assert isinstance(obj, B)
    function_that_needs_an_instance_of_b(obj)

This both keeps mypy happy and you sane, if you ever forget, what value you were expecting there.