Narrowing for complex targets (attribute expressions, subscripts)

We support type narrowing for attributes and subscripts.

Attribute narrowing

Basic

from ty_extensions import Unknown

class C:
    x: int | None = None

c = C()

reveal_type(c.x)  # revealed: int | None

if c.x is not None:
    reveal_type(c.x)  # revealed: int
else:
    reveal_type(c.x)  # revealed: None

if c.x is not None:
    c.x = None

reveal_type(c.x)  # revealed: None

c = C()

if c.x is None:
    c.x = 1

reveal_type(c.x)  # revealed: int

class _:
    reveal_type(c.x)  # revealed: int

c = C()

class _:
    if c.x is None:
        c.x = 1
    reveal_type(c.x)  # revealed: int

# TODO: should be `int`
reveal_type(c.x)  # revealed: int | None

class D:
    x = None

def unknown() -> Unknown:
    return 1

d = D()
reveal_type(d.x)  # revealed: None | Unknown
d.x = 1
reveal_type(d.x)  # revealed: Literal[1]
d.x = unknown()
reveal_type(d.x)  # revealed: Unknown

class E:
    x: int | None = None

e = E()

if e.x is not None:
    class _:
        reveal_type(e.x)  # revealed: int

Narrowing can be "reset" by assigning to the attribute:

c = C()

if c.x is None:
    reveal_type(c.x)  # revealed: None
    c.x = 1
    reveal_type(c.x)  # revealed: Literal[1]
    c.x = None
    reveal_type(c.x)  # revealed: None

reveal_type(c.x)  # revealed: int | None

Narrowing can also be "reset" by assigning to the object:

c = C()

if c.x is None:
    reveal_type(c.x)  # revealed: None
    c = C()
    reveal_type(c.x)  # revealed: int | None

reveal_type(c.x)  # revealed: int | None

Multiple predicates

class C:
    value: str | None

def foo(c: C):
    # The truthiness check `c.value` narrows to `str & ~AlwaysFalsy`.
    # The subsequent `len(c.value)` doesn't narrow further since `str` is not narrowable by len().
    if c.value and len(c.value):
        reveal_type(c.value)  # revealed: str & ~AlwaysFalsy

    # error: [invalid-argument-type] "Argument to function `len` is incorrect: Expected `Sized`, found `str | None`"
    if len(c.value) and c.value:
        reveal_type(c.value)  # revealed: str & ~AlwaysFalsy

    if c.value is None or not len(c.value):
        reveal_type(c.value)  # revealed: str | None
    else:  # c.value is not None and len(c.value)
        # `c.value is not None` narrows to `str`, but `str` is not narrowable by len().
        reveal_type(c.value)  # revealed: str

Generic class

toml

[environment]
python-version = "3.12"

class C[T]:
    x: T
    y: T

    def __init__(self, x: T):
        self.x = x
        self.y = x

def f(a: int | None):
    c = C(a)
    reveal_type(c.x)  # revealed: int | None
    reveal_type(c.y)  # revealed: int | None
    if c.x is not None:
        reveal_type(c.x)  # revealed: int
        # In this case, it may seem like we can narrow it down to `int`,
        # but different values may be reassigned to `x` and `y` in another place.
        reveal_type(c.y)  # revealed: int | None

def g[T](c: C[T]):
    reveal_type(c.x)  # revealed: T@g
    reveal_type(c.y)  # revealed: T@g
    reveal_type(c)  # revealed: C[T@g]

    if isinstance(c.x, int):
        reveal_type(c.x)  # revealed: T@g & int
        reveal_type(c.y)  # revealed: T@g
        reveal_type(c)  # revealed: C[T@g]
    if isinstance(c.x, int) and isinstance(c.y, int):
        reveal_type(c.x)  # revealed: T@g & int
        reveal_type(c.y)  # revealed: T@g & int
        # TODO: Probably better if inferred as `C[T & int]` (mypy and pyright don't support this)
        reveal_type(c)  # revealed: C[T@g]

With intermediate scopes

class C:
    def __init__(self):
        self.x: int | None = None
        self.y: int | None = None

c = C()
reveal_type(c.x)  # revealed: int | None
if c.x is not None:
    reveal_type(c.x)  # revealed: int
    reveal_type(c.y)  # revealed: int | None

if c.x is not None:
    def _():
        reveal_type(c.x)  # revealed: int | None

def _():
    if c.x is not None:
        reveal_type(c.x)  # revealed: int

Subscript narrowing

Number subscript

def _(t1: tuple[int | None, int | None], t2: tuple[int, int] | tuple[None, None]):
    if t1[0] is not None:
        reveal_type(t1[0])  # revealed: int
        reveal_type(t1[1])  # revealed: int | None

    n = 0
    if t1[n] is not None:
        # Narrowing the individual element type with a non-literal subscript is not supported
        reveal_type(t1[0])  # revealed: int | None
        reveal_type(t1[n])  # revealed: int | None
        reveal_type(t1[1])  # revealed: int | None

    # However, we can still discriminate between tuples in a union using a variable index:
    if t2[n] is not None:
        reveal_type(t2)  # revealed: tuple[int, int]

    if t2[0] is not None:
        reveal_type(t2)  # revealed: tuple[int, int]
        reveal_type(t2[0])  # revealed: int
        reveal_type(t2[1])  # revealed: int
    else:
        reveal_type(t2)  # revealed: tuple[None, None]
        reveal_type(t2[0])  # revealed: None
        reveal_type(t2[1])  # revealed: None

    if t2[0] is None:
        reveal_type(t2)  # revealed: tuple[None, None]
    else:
        reveal_type(t2)  # revealed: tuple[int, int]

def _(t3: tuple[int, str] | tuple[None, None] | tuple[bool, bytes]):
    # Narrow to tuples where first element is not None
    if t3[0] is not None:
        reveal_type(t3)  # revealed: tuple[int, str] | tuple[bool, bytes]

    # Narrow to tuples where first element is None
    if t3[0] is None:
        reveal_type(t3)  # revealed: tuple[None, None]

def _(t4: tuple[bool, int] | tuple[bool, str]):
    # Both tuples have bool at index 0, which is not disjoint from True,
    # so neither gets filtered out when checking `is True`
    if t4[0] is True:
        reveal_type(t4)  # revealed: tuple[bool, int] | tuple[bool, str]

def _(t5: tuple[int, None] | tuple[None, int]):
    # Narrow on second element (index 1)
    if t5[1] is not None:
        reveal_type(t5)  # revealed: tuple[None, int]
    else:
        reveal_type(t5)  # revealed: tuple[int, None]

    # Negative index
    if t5[-1] is None:
        reveal_type(t5)  # revealed: tuple[int, None]

def _(t6: tuple[int, ...] | tuple[None, None]):
    # Variadic tuple at index 0 has element type `int` (not a union),
    # so `tuple[None, None]` gets filtered out
    if t6[0] is not None:
        reveal_type(t6)  # revealed: tuple[int, ...]

def _(t6b: tuple[int, ...] | tuple[None, ...]):
    # Both variadic: `int` is disjoint from None, `None` is not disjoint from None
    if t6b[0] is not None:
        reveal_type(t6b)  # revealed: tuple[int, ...]
    else:
        reveal_type(t6b)  # revealed: tuple[None, ...]

def _(t7: tuple[int, int] | tuple[None, None]):
    # Index out of range for both tuples - no narrowing, but errors are emitted
    # error: [index-out-of-bounds] "Index 5 is out of bounds for tuple `tuple[int, int]` with length 2"
    # error: [index-out-of-bounds] "Index 5 is out of bounds for tuple `tuple[None, None]` with length 2"
    if t7[5] is not None:
        reveal_type(t7)  # revealed: tuple[int, int] | tuple[None, None]

def _(t8: tuple[int, int, int] | tuple[None, None]):
    # Index in range for first tuple but out of range for second
    # error: [index-out-of-bounds] "Index 2 is out of bounds for tuple `tuple[None, None]` with length 2"
    if t8[2] is not None:
        reveal_type(t8)  # revealed: tuple[int, int, int] | tuple[None, None]

def _(t9: tuple[int | None, str] | tuple[str, int]):
    # When the element type is a union (like `int | None`), we can't filter
    # out the tuple.
    if t9[0] is not None:
        reveal_type(t9)  # revealed: tuple[int | None, str] | tuple[str, int]

Tagged unions of tuples (equality narrowing)

Narrow unions of tuples based on literal tag elements using == comparison:

from typing import Literal

class A: ...
class B: ...
class C: ...

def _(x: tuple[Literal["tag1"], A] | tuple[Literal["tag2"], B, C]):
    if x[0] == "tag1":
        reveal_type(x)  # revealed: tuple[Literal["tag1"], A]
        reveal_type(x[1])  # revealed: A
    else:
        reveal_type(x)  # revealed: tuple[Literal["tag2"], B, C]
        reveal_type(x[1])  # revealed: B
        reveal_type(x[2])  # revealed: C

def _(x: tuple[Literal["tag1"], A] | tuple[Literal["tag2"], B, C]):
    if x[0] != "tag1":
        reveal_type(x)  # revealed: tuple[Literal["tag2"], B, C]
    else:
        reveal_type(x)  # revealed: tuple[Literal["tag1"], A]

# With int literals
def _(x: tuple[Literal[1], A] | tuple[Literal[2], B]):
    if x[0] == 1:
        reveal_type(x)  # revealed: tuple[Literal[1], A]
    else:
        reveal_type(x)  # revealed: tuple[Literal[2], B]

# With bytes literals
def _(x: tuple[Literal[b"a"], A] | tuple[Literal[b"b"], B]):
    if x[0] == b"a":
        reveal_type(x)  # revealed: tuple[Literal[b"a"], A]
    else:
        reveal_type(x)  # revealed: tuple[Literal[b"b"], B]

# Multiple tuple variants
def _(x: tuple[Literal["a"], A] | tuple[Literal["b"], B] | tuple[Literal["c"], C]):
    if x[0] == "a":
        reveal_type(x)  # revealed: tuple[Literal["a"], A]
    elif x[0] == "b":
        reveal_type(x)  # revealed: tuple[Literal["b"], B]
    else:
        reveal_type(x)  # revealed: tuple[Literal["c"], C]

# Using index 1 instead of 0
def _(x: tuple[A, Literal["tag1"]] | tuple[B, Literal["tag2"]]):
    if x[1] == "tag1":
        reveal_type(x)  # revealed: tuple[A, Literal["tag1"]]
    else:
        reveal_type(x)  # revealed: tuple[B, Literal["tag2"]]

# Works with reversed equality operands too.
def _(x: tuple[Literal["a"], A] | tuple[Literal["b"], B]):
    if "a" == x[0]:
        reveal_type(x)  # revealed: tuple[Literal["a"], A]
    else:
        reveal_type(x)  # revealed: tuple[Literal["b"], B]

# Works with reversed inequality operands too.
def _(x: tuple[Literal["a"], A] | tuple[Literal["b"], B]):
    if "a" != x[0]:
        reveal_type(x)  # revealed: tuple[Literal["b"], B]
    else:
        reveal_type(x)  # revealed: tuple[Literal["a"], A]

Narrowing is restricted to Literal tag elements. If any tuple has a non-literal type at the discriminating index, we can't safely narrow with equality:

def _(x: tuple[Literal["tag1"], A] | tuple[str, B]):
    # Can't narrow because second tuple has `str` (not literal) at index 0
    if x[0] == "tag1":
        reveal_type(x)  # revealed: tuple[Literal["tag1"], A] | tuple[str, B]
    else:
        # But we *can* narrow with inequality
        reveal_type(x)  # revealed: tuple[str, B]

If the index is out of bounds for any tuple in the union, we also skip narrowing (a diagnostic will be emitted elsewhere for the out-of-bounds access):

def _(x: tuple[A, Literal["a"]] | tuple[B]):
    # error: [index-out-of-bounds]
    if x[1] == "a":
        # Can't narrow because index 1 is out of bounds for second tuple
        reveal_type(x)  # revealed: tuple[A, Literal["a"]] | tuple[B]
    else:
        reveal_type(x)  # revealed: tuple[A, Literal["a"]] | tuple[B]

We can still narrow tuples when non-tuple types are present in the union:

def _(x: tuple[Literal["tag1"], A] | tuple[Literal["tag2"], B] | list[int]):
    if x[0] == "tag1":
        # A list of ints could have int subclasses in it,
        # and int subclasses could have custom `__eq__` methods such that they
        # compare equal to `"tag1"`, so `list[int]` cannot be narrowed out of this
        # union.
        reveal_type(x)  # revealed: tuple[Literal["tag1"], A] | list[int]

PEP 695 type aliases

Tuple narrowing also works when the union is defined via a PEP 695 type alias:

toml

[environment]
python-version = "3.12"

from typing import Literal

class A: ...
class B: ...

type TaggedTuple = tuple[Literal["a"], A] | tuple[Literal["b"], B]

def test_equality_narrowing(x: TaggedTuple):
    if x[0] == "a":
        reveal_type(x)  # revealed: tuple[Literal["a"], A]
    else:
        reveal_type(x)  # revealed: tuple[Literal["b"], B]

type NullableTuple = tuple[int, int] | tuple[None, None]

def test_is_narrowing(t: NullableTuple):
    if t[0] is not None:
        reveal_type(t)  # revealed: tuple[int, int]
    else:
        reveal_type(t)  # revealed: tuple[None, None]

# Nested type aliases (an alias referring to another alias) also work:
type InnerTagged = tuple[Literal["a"], A] | tuple[Literal["b"], B]
type OuterTagged = InnerTagged

def test_nested_equality_narrowing(x: OuterTagged):
    if x[0] == "a":
        reveal_type(x)  # revealed: tuple[Literal["a"], A]
    else:
        reveal_type(x)  # revealed: tuple[Literal["b"], B]

type InnerNullable = tuple[int, int] | tuple[None, None]
type OuterNullable = InnerNullable

def test_nested_is_narrowing(t: OuterNullable):
    if t[0] is not None:
        reveal_type(t)  # revealed: tuple[int, int]
    else:
        reveal_type(t)  # revealed: tuple[None, None]

String subscript

def _(d: dict[str, str | None]):
    if d["a"] is not None:
        reveal_type(d["a"])  # revealed: str
        reveal_type(d["b"])  # revealed: str | None

Combined attribute and subscript narrowing

class C:
    def __init__(self):
        self.x: tuple[int | None, int | None] = (None, None)

class D:
    def __init__(self):
        self.c: tuple[C] | None = None

d = D()
if d.c is not None and d.c[0].x[0] is not None:
    reveal_type(d.c[0].x[0])  # revealed: int

Narrowing with negative subscripts

Narrowing should work with negative subscripts like x[-1]:

def _(x: list[int | None]):
    if x[-1] is not None:
        reveal_type(x[-1])  # revealed: int

def _(x: list[str | None]):
    if x[-1] is None:
        reveal_type(x[-1])  # revealed: None
    else:
        reveal_type(x[-1])  # revealed: str

Nested negative subscripts should also work:

def _(x: list[list[int | None]]):
    if x[-1][-1] is not None:
        reveal_type(x[-1][-1])  # revealed: int

Mixed positive and negative subscripts:

def _(x: list[list[int | None]]):
    if x[0][-1] is not None:
        reveal_type(x[0][-1])  # revealed: int

    if x[-1][0] is not None:
        reveal_type(x[-1][0])  # revealed: int

Attribute access combined with negative subscripts:

class Container:
    items: list[int | None]

def _(c: Container):
    if c.items[-1] is not None:
        reveal_type(c.items[-1])  # revealed: int

Multiple conditions in an and chain:

def _(x: list[int | None]):
    # Narrowing should persist through `and` chains
    if x[-1] is not None and x[-1] > 0:
        reveal_type(x[-1])  # revealed: int

Negative indices with tuples:

def _(t: tuple[int, str, None] | tuple[None, None, int]):
    if t[-1] is not None:
        reveal_type(t)  # revealed: tuple[None, None, int]
    else:
        reveal_type(t)  # revealed: tuple[int, str, None]

    if t[-3] is not None:
        reveal_type(t)  # revealed: tuple[int, str, None]

Narrowing with explicit positive subscripts

Narrowing should work with explicit positive subscripts like x[+1]:

def _(x: list[int | None]):
    if x[+0] is not None:
        reveal_type(x[+0])  # revealed: int

    if x[+1] is not None:
        reveal_type(x[+1])  # revealed: int

Narrowing with boolean subscripts

Narrowing should work with boolean subscripts like x[True] and x[False]. We treat bool subscripts the same as int subscripts because True always has the same hash and index value as 1, and False always has the same hash and index value as 0:

def _(x: tuple[object, object]):
    if isinstance(x[True], str):
        reveal_type(x[True])  # revealed: str
        reveal_type(x[1])  # revealed: str

def _(x: list[int | None]):
    # x[True] is equivalent to x[1]
    if x[True] is not None:
        reveal_type(x[True])  # revealed: int

    # x[False] is equivalent to x[0]
    if x[False] is not None:
        reveal_type(x[False])  # revealed: int

Combined with other subscript types:

def _(x: list[list[int | None]]):
    if x[True][-1] is not None:
        reveal_type(x[True][-1])  # revealed: int

    if x[False][0] is not None:
        reveal_type(x[False][0])  # revealed: int

Narrowing with bytes literal subscripts

Narrowing should work with bytes literal subscripts like x[b"key"]:

def _(d: dict[bytes, str | None]):
    if d[b"key"] is not None:
        reveal_type(d[b"key"])  # revealed: str
        reveal_type(d[b"other"])  # revealed: str | None

Combined with attribute access:

class Container:
    data: dict[bytes, int | None]

def _(c: Container):
    if c.data[b"key"] is not None:
        reveal_type(c.data[b"key"])  # revealed: int