Singledispatch and type as an input argument

Question

I want to be able to do this:

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)

However, I get Invalid annotation for 'arg'. typing.Type[__main__.arg] is not a class.

I guess that singledispatch just does not support typing thoroughly for now. Is there any elegant workaround?

UPD I can not modify A and B.

Answer 1

You can give A a metaclass, and then all subclasses of A will be instances of that metaclass:

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