CODEWITHSUNDEEP
pythontyping
Maddy Guthridge
Maddy Guthridge

Reputation: 2001

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]:
    ...

Upvotes: 1

Views: 140

Answers (3)

Stefan
Stefan

Reputation: 776

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).

Upvotes: 3

Frost
Frost

Reputation: 11987

I am not sure this is possible.

The return type hint is supposed to tell you, without looking at the values of the arguments, what the function returns.

You are essentially implementing two different behaviors in this function:

  1. Return a Foo or crash
  2. Return a Foo or None

If you want to encode this in a type hint, you have to consider both cases. That means that the function is Optional[Foo].

If you want it not to be, I would suggest splitting this function into two different functions; one that raises and one that does dot:

def search_and_operate(f) -> Optional[Foo]:
    result = ... # some code that either results in a Foo object or None
    return result

def search_and_operate_or_raise(f) -> Foo:
    result = search_and_operate(f)
    if result is None:
        raise ValueError("No result")
    return result

Upvotes: 1

jfaccioni
jfaccioni

Reputation: 7539

Conceptually, I do not think that type hints that change depending on a parameter makes sense.

The type hint Optional[Foo] already encapsulates everything that searchAndOperate is able to return: anyone who calls it knows that it will return either Foo or None. Trying to "split" the return types based to the value of raise_on_failure couples the function's behavior to its implementation. What you should be doing is refactoring the function itself so that its behavior is non-ambiguous. The dissatisfaction with your current type hint stems from the fact that searchAndOperate tries to do too many things at once.

For example, searchAndOperate may be split into single-responsibility functions search and operate. Or you could just ditch the raise_on_failure parameter entirely, and leave to the caller to try/except the function call. If it becomes something you do a lot, you could wrap it as a separate function, something like:

def searchAndOperate(f) -> Foo:
    """Searches and returns Foo. Raises a ValueError if no result is found."""
    result = ...
    if result is None:
        raise ValueError("No result")
    return result

def trySearchAndOperate(f) -> Optional[Foo]:
    """Searches and returns Foo. Returns None if no result is found."""
    try:
        result = searchAndOperate(f)
    except ValueError:
        return None

This way, the caller already knows what to expect when no results are found. You might think that this is somewhat repetitive, but it is a common pattern: the only difference between 'abc'.index(s) and 'abc'.find(s) is that the first raises an exception if s isn't found in 'abc', and the second returns a default value of -1 in that case.

Upvotes: 2

Related Questions