CODEWITHSUNDEEP
pythonpython-3.xtypestyping
Chris J Harris
Chris J Harris

Reputation: 1851

Python type checking for functions

I'm having some issues figuring out how to mark up functions for type checking. Let's say I have the following situation:

class BaseClass:
    def __init__(self):
        self.name = 'base'


class MainClass(BaseClass):
    def __init__(self):
        self.mc_val = 1
        super().__init__()


def base_func(bc: BaseClass) -> BaseClass:
    bc.name = 'altered'
    return bc


def main_func(mc: MainClass) -> MainClass:
    mc.mc_val = 3
    return mc

mc = MainClass()
mc1 = base_func(mc)
mc2 = main_func(mc1)

This raises an issue in my type checker (Pycharm) because I've told it that mc1 is an instance of BaseClass, rather than MainClass, and so it doesn't fit the argument type for main_func.

So I need base_func to take an argument of type BaseClass (incl subclasses), but to return not BaseClass but the same type as was originally entered. I could have BaseClass both take and return, say, TypeVar('T') or similar, but then this rather defeats the point of type checking in the first place.

Any help greatly appreciated!

Upvotes: 2

Views: 81

Answers (1)

juanpa.arrivillaga
juanpa.arrivillaga

Reputation: 96257

You probably want a bounded type variable:

import typing
class BaseClass:
    def __init__(self):
        self.name = 'base'


class MainClass(BaseClass):
    def __init__(self):
        self.mc_val = 1
        super().__init__()

T = typing.TypeVar('T', bound=BaseClass)
def base_func(bc: T) -> T:
    bc.name = 'altered'
    return bc


def main_func(mc: MainClass) -> MainClass:
    mc.mc_val = 3
    return mc

mc = MainClass()
mc1 = base_func(mc)
mc2 = main_func(mc1)

Using mypy:

juan$ mypy test_typing.py
Success: no issues found in 1 source file

Upvotes: 4

Related Questions