How can I change the hinted return type of a function depending on the value of a parameter

I have a function where it usually returns an object that it searches for and performs some other actions. It raises an exception if it fails to find a match. Frequently, I don’t care if it finds a match or not, but not frequently enough that I’d consider removing the exception entirely. As such, the compromise I’ve made is to create a parameter raise_on_failure which defaults to True. If it is False, then None is returned rather than the reference.

def searchAndOperate(f, raise_on_failure: bool = True) -> Optional[Foo]:
    # Search
    result = ...
    if result is None:
        if raise_on_failure: raise ValueError("No results")
        else: return None
    # Operate
    ...
    # Then return the result
    return result

However, this has caused some issues with my type hinting. I want to make it so that if raise_on_failure is True, the function is guaranteed to return a Foo, such that it is only Optional[Foo] if raise_on_failure is False.

How can I do this? Something akin to the following snippet would be the most desirable, but I’m open to other ideas too.

def searchAndOperate(
    f,
    raise_on_failure: bool = True
) -> Foo if raise_on_failure else Optional[Foo]:
    ...

Answer

Others already mentioned that it might be better to refactor the function searchAndOperate into two separate functions.

Just for completeness however, it’s actually possible to create type annotations for the original behaviour by using the @overload decorator:

@overload
def searchAndOperate(f, raise_on_failure: Literal[True] = True) -> Foo: ...
@overload
def searchAndOperate(f, raise_on_failure: Literal[False]) -> Optional[Foo]: ...
@overload
def searchAndOperate(f, raise_on_failure: bool = True) -> Optional[Foo]: ...
def searchAndOperate(f, raise_on_failure = True):
    # Search
    result = ...
    if result is None:
        if raise_on_failure: raise ValueError("No results")
        else: return None
    # Operate
    ...
    # Then return the result
    return result

Note that there’s an overload for each combination plus one generic fallback. For more comprehensive function signatures that overload list can get quite long. A static type checker like mypy now infers the following return types:

a = searchAndOperate(f)  # Foo
b = searchAndOperate(f, raise_on_failure=True)  # Foo
c = searchAndOperate(f, raise_on_failure=False)  # Foo | None
d = searchAndOperate(f, raise_on_failure=(10 > 2))  # Foo | None
e = searchAndOperate(f, raise_on_failure=boolarg)  # Foo | None

When using no keyword argument or the literal arguments True or False, the type checker correctly infers one of the two specialised overloads. When passing a boolean expression or variable the type checker can infer the generic overload only (since it cannot know the actual value without running the code).