python - Singledispatch 和 type 作为输入参数

Question

我希望能够做到这一点:

from typing import Type
from functools import singledispatch

class A:
    pass

class B(A):
    pass

@singledispatch
def foo(arg):
    print('default')

@foo.register
def _(arg: Type[A]):
    print(arg)

foo(A)
foo(B)

但是，我得到了 “arg”的无效注释。 typing.Type[__main__.arg] 不是一个类。

我猜 singledispatch 目前还不能完全支持 typing。有什么优雅的解决方法吗？

UPD我无法修改 A 和 B。

Answer 1

您可以给 A 一个元类，然后 A 的所有子类都将是该元类的实例:

from functools import singledispatch

class AMeta(type): pass

class A(metaclass=AMeta): pass

class B(A): pass

@singledispatch
def foo(arg):
    print('default')

@foo.register
def _(arg: AMeta):
    print('A or B')

foo(A) # A or B
foo(B) # A or B
foo(A()) # default

class C: pass

foo(C) # default