[ty] Add precise iteration and unpacking inference for string literals and bytes literals (#20023)

## Summary

Previously we held off from doing this because we weren't sure that it
was worth the added complexity cost. But our code has changed in the
months since we made that initial decision, and I think the structure of
the code is such that it no longer really leads to much added complexity
to add precise inference when unpacking a string literal or a bytes
literal.

The improved inference we gain from this has real benefits to users (see
the mypy_primer report), and this PR doesn't appear to have a
performance impact.

## Test plan

mdtests
This commit is contained in:
Alex Waygood 2025-08-22 19:33:08 +01:00 committed by GitHub
parent 796819e7a0
commit bc6ea68733
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 203 additions and 43 deletions

View file

@ -74,6 +74,22 @@ def match_non_exhaustive(x: Literal[0, 1, "a"]):
# this diagnostic is correct: the inferred type of `x` is `Literal[1]`
assert_never(x) # error: [type-assertion-failure]
# This is based on real-world code:
# https://github.com/scipy/scipy/blob/99c0ef6af161a4d8157cae5276a20c30b7677c6f/scipy/linalg/tests/test_lapack.py#L147-L171
def exhaustiveness_using_containment_checks():
for norm_str in "Mm1OoIiFfEe":
if norm_str in "FfEe":
return
else:
if norm_str in "Mm":
return
elif norm_str in "1Oo":
return
elif norm_str in "Ii":
return
assert_never(norm_str)
```
## Checks on enum literals

View file

@ -755,6 +755,18 @@ def f(never: Never):
reveal_type(x) # revealed: Unknown
```
## Iterating over literals
```py
from typing import Literal
for char in "abcde":
reveal_type(char) # revealed: Literal["a", "b", "c", "d", "e"]
for char in b"abcde":
reveal_type(char) # revealed: Literal[97, 98, 99, 100, 101]
```
## A class literal is iterable if it inherits from `Any`
A class literal can be iterated over if it has `Any` or `Unknown` in its MRO, since the

View file

@ -523,8 +523,8 @@ def f(x: MixedTupleSubclass):
```py
a, b = "ab"
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
reveal_type(a) # revealed: Literal["a"]
reveal_type(b) # revealed: Literal["b"]
```
### Uneven unpacking (1)
@ -570,37 +570,37 @@ reveal_type(d) # revealed: Unknown
```py
(a, *b, c) = "ab"
reveal_type(a) # revealed: LiteralString
reveal_type(a) # revealed: Literal["a"]
reveal_type(b) # revealed: list[Never]
reveal_type(c) # revealed: LiteralString
reveal_type(c) # revealed: Literal["b"]
```
### Starred expression (3)
```py
(a, *b, c) = "abc"
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: list[LiteralString]
reveal_type(c) # revealed: LiteralString
reveal_type(a) # revealed: Literal["a"]
reveal_type(b) # revealed: list[Literal["b"]]
reveal_type(c) # revealed: Literal["c"]
```
### Starred expression (4)
```py
(a, *b, c, d) = "abcdef"
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: list[LiteralString]
reveal_type(c) # revealed: LiteralString
reveal_type(d) # revealed: LiteralString
reveal_type(a) # revealed: Literal["a"]
reveal_type(b) # revealed: list[Literal["b", "c", "d"]]
reveal_type(c) # revealed: Literal["e"]
reveal_type(d) # revealed: Literal["f"]
```
### Starred expression (5)
```py
(a, b, *c) = "abcd"
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
reveal_type(c) # revealed: list[LiteralString]
reveal_type(a) # revealed: Literal["a"]
reveal_type(b) # revealed: Literal["b"]
reveal_type(c) # revealed: list[Literal["c", "d"]]
```
### Starred expression (6)
@ -650,8 +650,114 @@ reveal_type(b) # revealed: Unknown
```py
(a, b) = "\ud800\udfff"
reveal_type(a) # revealed: Literal["<22>"]
reveal_type(b) # revealed: Literal["<22>"]
```
### Very long literal
```py
string = "very long stringgggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg"
a, *b = string
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
reveal_type(b) # revealed: list[LiteralString]
```
## Bytes
### Simple unpacking
```py
a, b = b"ab"
reveal_type(a) # revealed: Literal[97]
reveal_type(b) # revealed: Literal[98]
```
### Uneven unpacking (1)
```py
# error: [invalid-assignment] "Not enough values to unpack: Expected 3"
a, b, c = b"ab"
reveal_type(a) # revealed: Unknown
reveal_type(b) # revealed: Unknown
reveal_type(c) # revealed: Unknown
```
### Uneven unpacking (2)
```py
# error: [invalid-assignment] "Too many values to unpack: Expected 2"
a, b = b"abc"
reveal_type(a) # revealed: Unknown
reveal_type(b) # revealed: Unknown
```
### Starred expression (1)
```py
# error: [invalid-assignment] "Not enough values to unpack: Expected at least 3"
(a, *b, c, d) = b"ab"
reveal_type(a) # revealed: Unknown
reveal_type(b) # revealed: list[Unknown]
reveal_type(c) # revealed: Unknown
reveal_type(d) # revealed: Unknown
```
```py
# error: [invalid-assignment] "Not enough values to unpack: Expected at least 3"
(a, b, *c, d) = b"a"
reveal_type(a) # revealed: Unknown
reveal_type(b) # revealed: Unknown
reveal_type(c) # revealed: list[Unknown]
reveal_type(d) # revealed: Unknown
```
### Starred expression (2)
```py
(a, *b, c) = b"ab"
reveal_type(a) # revealed: Literal[97]
reveal_type(b) # revealed: list[Never]
reveal_type(c) # revealed: Literal[98]
```
### Starred expression (3)
```py
(a, *b, c) = b"abc"
reveal_type(a) # revealed: Literal[97]
reveal_type(b) # revealed: list[Literal[98]]
reveal_type(c) # revealed: Literal[99]
```
### Starred expression (4)
```py
(a, *b, c, d) = b"abcdef"
reveal_type(a) # revealed: Literal[97]
reveal_type(b) # revealed: list[Literal[98, 99, 100]]
reveal_type(c) # revealed: Literal[101]
reveal_type(d) # revealed: Literal[102]
```
### Starred expression (5)
```py
(a, b, *c) = b"abcd"
reveal_type(a) # revealed: Literal[97]
reveal_type(b) # revealed: Literal[98]
reveal_type(c) # revealed: list[Literal[99, 100]]
```
### Very long literal
```py
too_long = b"very long bytes stringggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg"
a, *b = too_long
reveal_type(a) # revealed: int
reveal_type(b) # revealed: list[int]
```
## Union
@ -714,7 +820,7 @@ def _(arg: tuple[int, tuple[str, bytes]] | tuple[tuple[int, bytes], Literal["ab"
a, (b, c) = arg
reveal_type(a) # revealed: int | tuple[int, bytes]
reveal_type(b) # revealed: str
reveal_type(c) # revealed: bytes | LiteralString
reveal_type(c) # revealed: bytes | Literal["b"]
```
### Starred expression
@ -785,8 +891,8 @@ from typing import Literal
def _(arg: tuple[int, int] | Literal["ab"]):
a, b = arg
reveal_type(a) # revealed: int | LiteralString
reveal_type(b) # revealed: int | LiteralString
reveal_type(a) # revealed: int | Literal["a"]
reveal_type(b) # revealed: int | Literal["b"]
```
### Custom iterator (1)