How to automatically infer type of child Class with mypy (without using cast)

Viewed 162

I have a generic class (Widget) and sub-classes with different attributes. I am using a dictionary with the key being an instance of these sub-classes in a function.

How to automatically infer with mpypy which sub-class I have?

More specifically, this is the code I have:

from dataclasses import dataclass
from typing import ClassVar, Dict


@dataclass
class Widget:
    """Generic class for widget"""


@dataclass
class Rectangle(Widget):
    """A Color Rectangle"""

    posx: int
    posy: int
    width: int = 500
    height: int = 200
    color: int = 0xFF0000


@dataclass
class Button(Widget):
    """A clickable button"""

    posx: int
    posy: int
    width: int = 200
    height: int = 100
    label: str = "some label"


@dataclass
class TestCase:
    ID: ClassVar[int]
    WIDGETS: ClassVar[Dict[str, Widget]]


class TestCaseButton(TestCase):
    ID = 1
    WIDGETS = {
        "ref": Button(posx=0, posy=0, label="Click me!"),
    }


class TestCaseRectangle(TestCase):
    ID = 2
    WIDGETS = {
        "ref": Rectangle(posx=0, posy=0, color=0xFFFFFF),
        "other": Button(posx=100, posy=100, label="DANGER!"),
    }


def run_tc_button(tc: TestCase):
    w = tc.WIDGETS["ref"]
    print(w.label)


def run_tc_rectangle(tc: TestCase):
    w = tc.WIDGETS["ref"]
    print(w.color)


if __name__ == "__main__":
    run_tc_button(TestCaseButton())
    run_tc_rectangle(TestCaseRectangle())

A check with mypy would return the following:

mypy_infer.py:55: error: "Widget" has no attribute "label"
        print(w.label)
              ^
mypy_infer.py:60: error: "Widget" has no attribute "color"
        print(w.color)
              ^
Found 2 errors in 1 file (checked 1 source file)

I know I could use cast like this to remove the errors:

from dataclasses import dataclass
from typing import ClassVar, Dict, cast

...

def run_tc_button(tc: TestCase):
    w = cast(Button, tc.WIDGETS["ref"])
    print(w.label)


def run_tc_rectangle(tc: TestCase):
    w = cast(Rectangle, tc.WIDGETS["ref"])
    print(w.color)

...

... but I find it a bit heavy to use cast each time in my functions.

How could I make the function automatically infer the type based on the values in the dictionary?

Note: the names of the keys are not fixed so a TypedDict will probably not solve the problem.

0 Answers
Related