[Python] Type annotation for a method returning an extended generic `TypedDict`

I come across the following code (simplified) ...

from typing import TypedDict
from typing import TypeVar


class BaseExternalState(TypedDict):
...


ExternalStateSchema = TypeVar("ExternalStateSchema", bound=BaseExternalState)


class Foo:
def __init__(self, es: type[ExternalStateSchema]) -> None:
self.es = es

class InternalState(es):
retry_start_time: float

self.internal_state = InternalState

@staticmethod
def get_internal_state(internal_state):
return internal_state

def main(self):
return self.get_internal_state(self.internal_state)



class Bar(BaseExternalState):
bar_param: int


class Baz(BaseExternalState):
baz_param: int


print(Foo(Bar).main()(retry_start_time=42, baz_param=42))
print(Foo(Baz).main()(retry_start_time=42, bar_param=42))


... and got stuck trying to come up with a type annotation for get_internal_state(). It is returning a type which inherits both ExternalStateSchema generic and InternalState. While InternalState part is easy - I can, e.g., introduce a top level protocol or something, I cannot wrap my mind around how to also define this type being a child of a generic ExternalStateSchema.

Any ideas on how to type annotate this code completely? Or is it a lost cause and I need to change the implementation to be able to annotate it properly?

Continue reading...