How to typehint np.nan, value input and function return

Question

For example, if i have:

def f(x):
    return x

Where x is expected to be a NaN value - how's that supposed to be typehinted?

If i hint with np.nan i get the error: error: Variable "numpy.nan" is not valid as a type. I could use np.float64, but this feels too broad for what i want here.

eg:

def f(x : np.nan) -> np.nan:
    return x

Answer 1

In order to specify values (not types) you have to use Literal types.

However, Literal types may contain one or more literal bools, ints, strs, bytes, and enum values.

As long as floats are excluded you need to create an enum defining the desired value:

from enum import Enum
from typing import Literal

class MyValues(Enum):
    nan: np.nan
    zero: 0.
    i_like_seven: 7.
    ...

NaN = Literal[MyValues.nan]    # A convenience type alias

Then you can use (hopefully in a more useful function):

def f(x: NaN) -> NaN:
    return x

If you check it using mypy:

>>> f(3.)       # mypy complains: Expected type 'MyValues', got 'float' instead
>>> f(np.nan)   # mypy is happy