如何根据传入实例的方法返回值在python中键入函数？

Question

在 TypeScript 中，我可以定义一个函数，使得返回值根据传递给函数的实例的行为而不同。 （见下面的例子。）

如何在 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

Answer 1

你有几个选择：

使用

@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"))

返回类型取决于传入的内容：

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

使用泛型

from typing import TypeVar

T = TypeVar('T')

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

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

另一种选择：

可以使用

Union[]

类型，但这会使您的函数也返回联合类型，而不是实际的运行时类型。

Answer 2

您可以对函数使用“singledispatch”，对类方法使用“singledispatchmethod”。

此外，此代码通过了 mypy 检查，因此代码的输入注释没问题。

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

Answer 3

这似乎有效。奇怪的是我没有明确定义返回类型。但是，如果我只是将其保留为未定义，则类型提示器似乎可以解决问题：

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())