Type narrowing by a generic function in Python

253 Views Asked by At

I need to narrow type of a class attribute. It is typed as MyType | None and it is expected not be None in a specific part of the code. I am using this attribute in a generator expression so I cannot directly use a condition with raise or assert to narrow it.

Is the code below the correct way to define a generic type-narrowing function? I.e. define argument as a generic type in union with None but the return type just as the generic type.

from typing import TypeVar

ValueT = TypeVar('ValueT')

def fail_if_none(value: ValueT | None) -> ValueT:
    """Raises an exception if the value is None.

    Args:
        value: The value to check.

    Raises:
        ValueError: If the value is None.

    Returns:
        The value, guaranteed to not be None.
    """
    if value is None:
        raise ValueError("Value cannot be None.")
    return value

Does the code guarantee that the return type ValueT will always be seen as not including None? Is not there a better way?


Here is an example showing how to use the type-narrowing function:

def takes_int(value: int) -> int:
    """Take an int and return it."""
    return value

sequence: list[int | None] = [1, 2, 3, 4, None, 5]

values1 = (takes_int(value) for value in sequence)
# error: Argument 1 to "takes_int" has incompatible type "int | None";
# expected "int"  [arg-type]

values2 = (takes_int(fail_if_none(value)) for value in sequence)
# OK, no type-checker errors

I have tested the code with mypy and Pyright. In both type-checkers it seems to work the way I wanted.

0

There are 0 best solutions below