Support type narrowing literals using typing.get_args()

[SRv6d]

Feature

When using typing.get_args() to ensure that a given string is of a literal type, I would expect mypy to behave similarly than when isinstance() checks are used.

Pitch
Take the following code:

from typing import Literal, TypeAlias, get_args

ExpectedUserInput: TypeAlias = Literal[
    "these", "strings", "are", "expected", "user", "input"
]


def external_function(input: str) -> str:
    if input not in get_args(ExpectedUserInput):
        raise ValueError("Invalid input.")
    return _internal_function(input)


def _internal_function(input: ExpectedUserInput) -> str:
    return f"User input: {input}"

Mypy does not detect that the string must be of the correct type, since the code path is unreachable for an invalid value and outputs an error, requiring the use of casts even though the value is guaranteed to be of the correct type:

...: error: Argument 1 to "_internal_function" has incompatible type "str"; expected "Literal['these', 'strings', 'are', 'expected', 'user', 'input']"  [arg-type]

[ckp95]

+1 to this. I had wanted to use Literal as a kind of witness that untrusted input has been sanitized, but it's impractical without a way to automatically check against all possible members of the union.

You can do this:

from typing import Literal, Optional

ValidToken = Literal["foo", "bar", "baz"]

def parse_token(s: str) -> Optional[ValidToken]:
    if s == "foo" or s == "bar" or s == "baz":
        return s
    else:
        return None

if __name__ == "__main__":
    user_input = input()
    token = parse_token(user_input)  # token has type Optional[ValidToken]

But this requires laboriously writing out the members of the union in duplicate. And you have to update it in both places when you add or remove members. If we could use if s in get_args(ValidToken) to narrow the type, that would open a lot of possibilities for type-driven parsing and domain modelling.

Would be even cooler if it could somehow be made to work like if isinstance(s, ValidToken) but I think that would be an abuse of notation.

[Jordandev678]

I've just hit the desire to do this as well. I have a little time and I'm up for taking a run at and see if it's not too hard to do.
It seems to need two things.

1. Ability to narrow strings to Union[Literal[], Literal[], ...] with "in"
2. Handling for get_args() to allow mypy to understand get_args(Union[Literal['x'], Literal['y']]) returns ('x', 'y') / Tuple[Union[Literal['x'], Literal['y']]]

I took a quick look at microsoft/pyright#5837 which is similar. Based on that pyright already supports the narrowing part. Though there wasn't much support for adding the get_args() handling as it mixes type and value expressions.
I get that point, but it is a standard function, the behavior for Literal["x", "y"] returning ('x', 'y') is defined, and it's a natural way to want to write a runtime type guard that in my opinion a static checker should be able to identify as a type guard (automatically) in the case of strings/literals.

typing.get_args(tp)¶
Get type arguments with all substitutions performed: for a typing object of the form X[Y, Z, ...] return (Y, Z, ...).

Assuming the type narrowing isn't objectionable it seems the main question is how to handle get_args().
I'd propose if it's given a type that resolves to Union[Literal["a"], Literal["b"]] have it's return type set to Tuple[Union[Literal["a"], Literal["b"]]]. Otherwise leave it as is which seems to be Tuple[Any].

Thoughts?

[JelleZijlstra]

That makes sense. Your (1) has a number of open issues about it already (I think #3229 is the oldest) and should be implemented independently.

Once that is done, we can look into special handling for get_args() to fix this issue.