False-positive type error using TypeVar with "is None" conditional

[tpajenkamp-dspace]

Describe the bug Pyright reports an error for functions like the one below which uses a TypeVar for both an argument and its return type where None is one of the constraints. The type of the type variable is not properly narrowed (or whatever the correct term is in this case) by the arg is None conditional. If I use isinstance(arg, types.NoneType) instead it works, but this form does not seem idiomatic to me.

Code

import typing

OptionalStr = typing.TypeVar("OptionalStr", str, None)

def string_fun(arg: OptionalStr) -> OptionalStr:
    typing.reveal_type(arg)     # Type of "arg" is "OptionalStr@string_fun"
    if arg is None:
        typing.reveal_type(arg) # Type of "arg" is "None*"
        return None             # error: Type "None" is not assignable to type "OptionalStr@string_fun" (reportReturnType)

    typing.reveal_type(arg)     # Type of "arg" is "str*"
    return f"s: {arg}"          # error: Type "str" is not assignable to type "OptionalStr@string_fun" (reportReturnType)

VS Code extension or command-line pyright 1.1.388 & 'C:\Python312\python.exe' -m pyright .

[erictraut]

Pyright is working as intended here.

This is a limitation in how pyright implements value-constrained type variables. See this documentation for details.

I recommend avoiding value-constrained type variables. They are not well specified in the typing spec and have many limitations. They are not found in any other programming language — for good reason.

The recommended approach for your use case is to use an overload.

from typing import overload

@overload
def string_fun(arg: str) -> str: ...
@overload
def string_fun(arg: None) -> None: ...
def string_fun(arg: str | None) -> str | None:
    if arg is None:
        return None
    return f"s: {arg}"