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

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?"

CodePudding user response：

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 Generic, 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.