Type hint issue when base class is parameterized on a value type and has methods that return that type

This a toy example to illustrate the issue:

from typing import Type
from dataclasses import dataclass, asdict
import json

@dataclass
class ValueType1:
    x: int
    y: int

@dataclass
class ValueType2:
    x: int
    z: str

class FooBase:
    def __init__(self, value_cls: Type[ValueType1|ValueType2], name: str):
        self.name = name
        self.value_cls = value_cls
    def save(self, value: ValueType1|ValueType2):
        with open(self.name+'.json', 'w') as f:
            json.dump(asdict(value), f)
    def load(self) -> ValueType1|ValueType2:
        with open(self.name+'.json', 'r') as f:
            return self.value_cls(**json.load(f))
    
class Foo1(FooBase):
    def __init__(self):
        super().__init__(value_cls=ValueType1, name='1')
class Foo2(FooBase):
    def __init__(self):
        super().__init__(value_cls=ValueType2, name='2')

foo1 = Foo1()
foo2 = Foo2()
foo1.save(ValueType1(x=10, y=20))
foo2.save(ValueType2(x=10, z='a'))
res1 = foo1.load()
res2 = foo2.load()
print(res1.y)
print(res2.z)

This works fine and prints 20 and ‘a’ as expected but typing errors are shown at res1.y and res2.z of this sort (only res1.y errors shown):

Type of "y" is unknownPylancereportUnknownMemberType
Type of "y" is partially unknown
  Type of "y" is "int | Unknown"PylancereportUnknownMemberType
Argument type is partially unknown
  Argument corresponds to parameter "values" in function "print"
  Argument type is "int | Unknown"PylancereportUnknownArgumentType
Cannot access member "y" for type "ValueType2"
  Member "y" is unknownPylancereportGeneralTypeIssues
(variable) y: int | Unknown

I’m wondering if there is some type hinting magic that can resolve this and make it clear that res1 is of type ValueType1. Thanks!

>Solution :

you could do this:

from typing import Callable

class Foo1(FooBase):
    load: Callable[[], ValueType1]
    def __init__(self):
        super().__init__(value_cls=ValueType1, name='1')

class Foo2(FooBase):
    load: Callable[[], ValueType2]
    def __init__(self):
        super().__init__(value_cls=ValueType2, name='2')

Some editors will now pick load() up as a attribute and not a function and color it wrong though. But type hints will work properly