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.