Search code examples
pythonpycharmpython-typing

python type hint, return same type as input


Well I'm considering a case where I have a baseclass with (several) children. I have a function that takes a list of baseclass objects, and returns a new list with these objects in it.

Now if I would use a child class obviously the return is a list of those child class objects: consider the following trivial case:

from typing import Sequence, List, TypeVar


class BaseClass:
    def __init__(self, data=None, *args, **kwargs):
        super().__init__()
        self.CanCalculate = True
        if data is None:
            data = []
            self.CanCalculate = False
        self._mydata = list(data)
        self.multiplier = 1

    @property
    def data(self):
        return self._mydata


class ChildClass(BaseClass):
    def sum_mul_data(self):
        return self.multiplier * sum(self.data)


class SecondChildClass(BaseClass):
    def sum_div_data(self):
        return sum(self.data) / self.multiplier


def myFun(input: Sequence[BaseClass]) -> List[BaseClass]:
    out = []
    for n, i in enumerate(input):
        if i.CanCalculate:
            i.multiplier = 10**n
            out.append(i)
    return out


childs = [ChildClass([1,2,3,4]), ChildClass(), ChildClass([1,2,3,4])]

t = myFun(childs)
for idx in t:
    print(idx.sum_mul_data())


childs = [SecondChildClass([1,2,3,4]), SecondChildClass(), SecondChildClass([1,2,3,4])]

t = myFun(childs)
for idx in t:
    print(idx.sum_div_data())

Legal code: however pycharm (and the standard type hints) shows an error during static code analysis: "unresolved attribute reference" @ idx.sum_mul_data()

Now apparently this is due to pycharm thinking the return of the function is of type "BaseClass" - not a child. So how would I state: "return same type as input"?

I tried using a typevar: T = TypeVar("T", BaseClass), though that gave an actual error, that a single contraint can't be used in TypeVar. Interestingly using T = TypeVar("T", BaseClass, ChildClass) did work, and pycharm correctly deduced the type(hint) for sum_div_data.


Solution

  • You should use typevars with an upper bound: do T = TypeVar('T', bound=BaseClass) instead of T = TypeVar('T', BaseClass).

    Details:

    • When you do something like T = TypeVar('T', ClassA, ClassB, ClassC...), you're creating a type variable with a value restriction. That is, you're insisting that T must be exactly one of the classes you listed.

      This is why doing T = TypeVar('T', ClassA) is disallowed: the typevar can only be equal to one class, so you might as well just use the ClassA type directly.

    • When you do something like T = TypeVar('T', bound=ClassA), you're creating a type variable with an upper bound. You're insisting that T must either be ClassA, or any of its subclasses.