CODEWITHSUNDEEP
pythongenericspython-typingmypypython-3.11
Mobious
Mobious

Reputation: 122

In Python, how would you properly annotate a function to show the same return type as a function of a class member?

I'm using Python 3.11 with mypy as my type checker. Here's an example of what I'm trying to do.

from dataclasses import dataclass
from abc import ABC, abstractmethod
from typing import Any, TypeVar, Generic

class MyBase(ABC):
    @abstractmethod
    def my_func(self, *args, **kwargs) -> Any:
        raise NotImplementedError

class Bar(MyBase):
    def my_func(self, *args, **kwargs) -> int:
        return 5

class Baz(MyBase):
    def my_func(self, *args, **kwargs) -> str:
        return 'Hello'

T = TypeVar('T', bound=MyBase)

@dataclass
class Foo(Generic[T]):
    value: T

    def foo_call(self) -> ???:  # What do I put here? I want something like T.my_func.__annotations__['return']
        return self.value.my_func()

myFoo: Foo[Bar] = Foo(value=Bar())
myFoo.foo_call()  # this should give a return type-hint of int

I could just declare myFoo as a Foo[int], with foo_call(self) -> T, but then that wouldn't properly update if any changes are made to the return type of Bar.my_func. What's the proper way to do this?

EDIT

user2357112's method is helpful, but still requires updating the type hint in multiple places when the return type of my_func changes. It also requires the use of multiple TypeVars if I want to support multiple functions. For example:

from dataclasses import dataclass
from abc import ABC, abstractmethod
from typing import TypeVar, Generic

# need to add an additional TypeVar for every method I want to support
T = TypeVar('T')
TOther = TypeVar('TOther')

class MyBase(ABC, Generic[T, TOther]):
    @abstractmethod
    def my_func(self, *args, **kwargs) -> T:
        raise NotImplementedError

    @abstractmethod
    def other_func(self, *args, **kwargs) -> TOther:
        raise NotImplementedError

class Bar(MyBase[int, float]):
    def my_func(self, *args, **kwargs) -> int:
        return 5

    def other_func(self, *args, **kwargs) -> float:
        return 6.2


@dataclass
class Foo(Generic[T, TOther]):
    value: MyBase[T, TOther]

    def foo_call(self) -> T:
        return self.value.my_func()

    def foo_other(self) -> TOther:
        return self.value.other_func()

myFoo: Foo[int, float] = Foo(value=Bar())
myFoo.foo_call()
myFoo.foo_other()

I'd ideally like to be able to just pass the type of MyBase as the generic type hint for myFoo (e.g. myFoo: Foo[Bar] = ...) and then have the return types of methods like foo_call and foo_other be inferred from that.

Upvotes: 0

Views: 240

Answers (1)

user2357112
user2357112

Reputation: 280181

Most likely, MyBase should be generic, and Foo's parameter should be the my_func return type, not the MyBase subclass:

from dataclasses import dataclass
from abc import ABC, abstractmethod
from typing import TypeVar, Generic

T = TypeVar('T')

class MyBase(ABC, Generic[T]):
    @abstractmethod
    def my_func(self, *args, **kwargs) -> T:
        raise NotImplementedError

class Bar(MyBase[int]):
    def my_func(self, *args, **kwargs) -> int:
        return 5

class Baz(MyBase[str]):
    def my_func(self, *args, **kwargs) -> str:
        return 'Hello'

@dataclass
class Foo(Generic[T]):
    value: MyBase[T]

    def foo_call(self) -> T:
        return self.value.my_func()

Upvotes: 2

Related Questions