CODEWITHSUNDEEP
pythonpython-typingmypy
Roman Imankulov
Roman Imankulov

Reputation: 8817

How to type-annotate a Python function to mark it as returning None only when the input is None?

There is a typical pattern for data conversion in our code: when the value is None, we let it pass through. E.g.,

def capitalize(value):  
    if value is None:  
        return None  
    return value.capitalize()
     
# usage example:
assert capitalize(None) is None
assert capitalize("hello world") == "Hello world"

I can annotate it like this:

from typing import Optional


def capitalize(value: Optional[str]) -> Optional[str]:  
    if value is None:  
        return None  
    return value.capitalize()

Looks OK, but the following code

capitalize("Hello world").split()

will always make mypy complain.

error: Item "None" of "Optional[str]" has no attribute "split"

Is there a way to express with type annotations the conversion rule "None is always converted to None, and str is always converted to str?"

Upvotes: 2

Views: 750

Answers (1)

Laurent S
Laurent S

Reputation: 4326

This sounds like a use case for Generics with a value restriction.

The code below basically says T can be either a str or None, and the function definition says "this function returns the same type as what came in".

from typing import TypeVar

T = TypeVar("T", str, None)

def capitalize(value: T) -> T:
    if value is None:
        return None
    return value.capitalize()

capitalize("Hello world").split()

Running mypy on the above code seems to work fine, and:

capitalize(None).split()

causes mypy to complain: error: "None" has no attribute "split" which I think is what you're after.

Upvotes: 7

Related Questions