Just use a generic and make it bound to (A, B):

    from typing import *
    
    
    class _A:
        pass
    
    class _B(_A):
        pass
    
    A = NewType("A", _A)
    B = NewType("B", _B)
    
    def foo[T: (A, B)](val: T) -> T:
        return val
    
    a = A(_A())
    b = B(_B())
    
    _a = foo(a)
    _b = foo(b)
    
    reveal_type(_a)
    reveal_type(_b)

Playground here: https://mypy-play.net/?mypy=latest&python=3.12&gist=36573363...