How do type a function in python based on the return value of a method on an instance passed in?

In TypeScript, I can define a function such that the return value is different depending on the behavior of an instance that was passed into the function. (See the example below.)

How can I do something similar in Python?

class DocumentSuperClass {
   

}
class DocumentA extends DocumentSuperClass {}
class DocumentB extends DocumentSuperClass {}
// ... many more sibling classes

class SuperClass {
    doc: DocumentSuperClass

    meth() {
        return this.doc;
    }
}

class A {
    doc: DocumentA

    meth() {
        return this.doc;
    }
}

class B {
    doc: DocumentB

    meth() {
        return this.doc;
    }
}
// ... many more sibling classes

function a<T extends SuperClass>(input: T) {
    return input.meth() as ReturnType<T['meth']>;
}

const instanceOfDocumentA = a(new A());
const instanceOfDocumentB = a(new B());
// ... same behavior for many more sibling classes

You've got a few options:

Use @overload

from dataclasses import dataclass
from typing import overload

@dataclass
class A:
    val: int

@dataclass
class B:
    val: str


@overload
def a(arg: A) -> int: ...
@overload
def a(arg: B) -> str: ...

def a(arg: A | B):
    return arg.val


typed_as_int = a(A(123))
typed_as_str = a(B("abc"))

The return types depend on what's passed in:

Foo().one()         # -> Tuple[Any, ...] | None
Foo().one('tuple')  # -> Tuple[Any, ...] | None
Foo().one('dict')   # -> Dict[str, Any] | None

Use generics

from typing import TypeVar

T = TypeVar('T')

def foo(arg: T) -> T:
    return arg

foo(1)      # -> int
foo('bar')  # -> str

another alternative:

Could use Union[] type, but that makes your function return the union type as well, not the actual runtime type.

You can use 'singledispatch' for functions or 'singledispatchmethod' for class methods.

Also, this code passes mypy checks, so code's typing annotations are fine.

from functools import singledispatch

class A:
    def meth(self):
        return 'a'
    
class B:
    def meth(self):
        return 1
    
C = A | B

@singledispatch
def function(input_: C):
    pass
@function.register
def _(input_: A) -> str:
    return input_.meth()
@function.register
def _(input_: B) -> int:
    return input_.meth()

print(function(A()))
print(function(B()))

a
1

This seemed to work. It's odd to me that I'm not defining the return type explicitly. But if I just leave it undefined, the type hinter seems to figure things out:

class DocumentSuperClass:
    pass

class DocumentA(DocumentSuperClass):
    pass

class DocumentB(DocumentSuperClass):
    pass

class SuperClass:
    def __init__(self):
        self.doc = DocumentSuperClass

    def meth(self):
        return self.doc

class ClassA(SuperClass):
    def __init__(self):
        self.doc = DocumentA()

    def meth(self):
        return self.doc

class ClassB(SuperClass):
    def __init__(self):
        self.doc = DocumentB()

    def meth(self):
        return self.doc

def func(input:SuperClass):
    return input.meth()


instanceOfDocumentA = func(ClassA())
instanceOfDocumentB = func(ClassB())