function returns Union[x,y] but it is expected to be 'x' always --- How to prevent 'x' is not compatible with 'y'

[paxcodes]

Given something like this:

From third-party lib

from typing import Union

class A:
    pass

class B:
    pass

def some_third_party_func(param: str) -> Union[A,B]:
    if ("some_value"):
        return A()
    else:
        return B()

My own function

def my_func() -> A:
    a = some_third_party_func("some_value") # I expect this to return 'A()' all the time
    return a # Error: Expression of type A|B cannot be assigned to return type A

How can I work around the error?

I tried something like:

def my_func() -> A:
    a = some_third_party_func("some_value") # I expect this to return 'A()' all the time
    if type(a) == A:
        return a # Still gives the same error
    raise Exception("expecting a to be A")

Any advice?

Full-context The library I'm working with is bs4. The function I am calling is .find_next_sibling() which returns either a Tag or NavigableString --- but my own function typehints that it returns a Tag because that is what is expected --- If it is not a Tag then it raises an exception. But I'm having trouble with the type-hinting, thus this discussion.

[erictraut]

If the function's return type depends on the type of one of the input parameters, the library author should provide @overload functions to differentiate them.

You have a few workarounds:

def my_func() -> A:
    a = some_third_party_func("some_value")
    assert isinstance(a, A)
    return a

or

def my_func() -> A:
    a = some_third_party_func("some_value")
    if isinstance(a, A):
        return a
    raise Exception("expecting a to be A")

or

def my_func() -> A:
    a = some_third_party_func("some_value")
    return cast(A, a)

I recommend the first option since it is safe but still relatively terse.

Note that isinstance(a, A) is not the same as type(a) == A. The function claims that it returns type Union[A, B] which means that it returns either an instance of A, B, or a subtype of either. It doesn't necessarily return A. The expression isinstance(a, A) checks for A or any subtype of A.