Files
ruff/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md
Shunsuke Shibayama 0858896bc4 [ty] type narrowing by attribute/subscript assignments (#18041)
## Summary

This PR partially solves https://github.com/astral-sh/ty/issues/164
(derived from #17643).

Currently, the definitions we manage are limited to those for simple
name (symbol) targets, but we expand this to track definitions for
attribute and subscript targets as well.

This was originally planned as part of the work in #17643, but the
changes are significant, so I made it a separate PR.
After merging this PR, I will reflect this changes in #17643.

There is still some incomplete work remaining, but the basic features
have been implemented, so I am publishing it as a draft PR.
Here is the TODO list (there may be more to come):
* [x] Complete rewrite and refactoring of documentation (removing
`Symbol` and replacing it with `Place`)
* [x] More thorough testing
* [x] Consolidation of duplicated code (maybe we can consolidate the
handling related to name, attribute, and subscript)

This PR replaces the current `Symbol` API with the `Place` API, which is
a concept that includes attributes and subscripts (the term is borrowed
from Rust).

## Test Plan

`mdtest/narrow/assignment.md` is added.

---------

Co-authored-by: David Peter <sharkdp@users.noreply.github.com>
Co-authored-by: Carl Meyer <carl@astral.sh>
2025-06-04 17:24:27 -07:00

7.6 KiB

Narrowing by assignment

Attribute

Basic

class A:
    x: int | None = None
    y = None

    def __init__(self):
        self.z = None

a = A()
a.x = 0
a.y = 0
a.z = 0

reveal_type(a.x)  # revealed: Literal[0]
reveal_type(a.y)  # revealed: Literal[0]
reveal_type(a.z)  # revealed: Literal[0]

# Make sure that we infer the narrowed type for eager
# scopes (class, comprehension) and the non-narrowed
# public type for lazy scopes (function)
class _:
    reveal_type(a.x)  # revealed: Literal[0]
    reveal_type(a.y)  # revealed: Literal[0]
    reveal_type(a.z)  # revealed: Literal[0]

[reveal_type(a.x) for _ in range(1)]  # revealed: Literal[0]
[reveal_type(a.y) for _ in range(1)]  # revealed: Literal[0]
[reveal_type(a.z) for _ in range(1)]  # revealed: Literal[0]

def _():
    reveal_type(a.x)  # revealed: Unknown | int | None
    reveal_type(a.y)  # revealed: Unknown | None
    reveal_type(a.z)  # revealed: Unknown | None

if False:
    a = A()
reveal_type(a.x)  # revealed: Literal[0]
reveal_type(a.y)  # revealed: Literal[0]
reveal_type(a.z)  # revealed: Literal[0]

if True:
    a = A()
reveal_type(a.x)  # revealed: int | None
reveal_type(a.y)  # revealed: Unknown | None
reveal_type(a.z)  # revealed: Unknown | None

a.x = 0
a.y = 0
a.z = 0
reveal_type(a.x)  # revealed: Literal[0]
reveal_type(a.y)  # revealed: Literal[0]
reveal_type(a.z)  # revealed: Literal[0]

class _:
    a = A()
    reveal_type(a.x)  # revealed: int | None
    reveal_type(a.y)  # revealed: Unknown | None
    reveal_type(a.z)  # revealed: Unknown | None

def cond() -> bool:
    return True

class _:
    if False:
        a = A()
    reveal_type(a.x)  # revealed: Literal[0]
    reveal_type(a.y)  # revealed: Literal[0]
    reveal_type(a.z)  # revealed: Literal[0]

    if cond():
        a = A()
    reveal_type(a.x)  # revealed: int | None
    reveal_type(a.y)  # revealed: Unknown | None
    reveal_type(a.z)  # revealed: Unknown | None

class _:
    a = A()

    class Inner:
        reveal_type(a.x)  # revealed: int | None
        reveal_type(a.y)  # revealed: Unknown | None
        reveal_type(a.z)  # revealed: Unknown | None

# error: [unresolved-reference]
does.nt.exist = 0
# error: [unresolved-reference]
reveal_type(does.nt.exist)  # revealed: Unknown

Narrowing chain

class D: ...

class C:
    d: D | None = None

class B:
    c1: C | None = None
    c2: C | None = None

class A:
    b: B | None = None

a = A()
a.b = B()
a.b.c1 = C()
a.b.c2 = C()
a.b.c1.d = D()
a.b.c2.d = D()
reveal_type(a.b)  # revealed: B
reveal_type(a.b.c1)  # revealed: C
reveal_type(a.b.c1.d)  # revealed: D

a.b.c1 = C()
reveal_type(a.b)  # revealed: B
reveal_type(a.b.c1)  # revealed: C
reveal_type(a.b.c1.d)  # revealed: D | None
reveal_type(a.b.c2.d)  # revealed: D

a.b.c1.d = D()
a.b = B()
reveal_type(a.b)  # revealed: B
reveal_type(a.b.c1)  # revealed: C | None
reveal_type(a.b.c2)  # revealed: C | None
# error: [possibly-unbound-attribute]
reveal_type(a.b.c1.d)  # revealed: D | None
# error: [possibly-unbound-attribute]
reveal_type(a.b.c2.d)  # revealed: D | None

Do not narrow the type of a property by assignment

class C:
    def __init__(self):
        self._x: int = 0

    @property
    def x(self) -> int:
        return self._x

    @x.setter
    def x(self, value: int) -> None:
        self._x = abs(value)

c = C()
c.x = -1
# Don't infer `c.x` to be `Literal[-1]`
reveal_type(c.x)  # revealed: int

Do not narrow the type of a descriptor by assignment

class Descriptor:
    def __get__(self, instance: object, owner: type) -> int:
        return 1

    def __set__(self, instance: object, value: int) -> None:
        pass

class C:
    desc: Descriptor = Descriptor()

c = C()
c.desc = -1
# Don't infer `c.desc` to be `Literal[-1]`
reveal_type(c.desc)  # revealed: int

Subscript

Specialization for builtin types

Type narrowing based on assignment to a subscript expression is generally unsound, because arbitrary __getitem__/__setitem__ methods on a class do not necessarily guarantee that the passed-in value for __setitem__ is stored and can be retrieved unmodified via __getitem__. Therefore, we currently only perform assignment-based narrowing on a few built-in classes (list, dict, bytesarray, TypedDict and collections types) where we are confident that this kind of narrowing can be performed soundly. This is the same approach as pyright.

from typing import TypedDict
from collections import ChainMap, defaultdict

l: list[int | None] = [None]
l[0] = 0
d: dict[int, int] = {1: 1}
d[0] = 0
b: bytearray = bytearray(b"abc")
b[0] = 0
dd: defaultdict[int, int] = defaultdict(int)
dd[0] = 0
cm: ChainMap[int, int] = ChainMap({1: 1}, {0: 0})
cm[0] = 0
# TODO: should be ChainMap[int, int]
reveal_type(cm)  # revealed: ChainMap[Unknown, Unknown]

reveal_type(l[0])  # revealed: Literal[0]
reveal_type(d[0])  # revealed: Literal[0]
reveal_type(b[0])  # revealed: Literal[0]
reveal_type(dd[0])  # revealed: Literal[0]
# TODO: should be Literal[0]
reveal_type(cm[0])  # revealed: Unknown

class C:
    reveal_type(l[0])  # revealed: Literal[0]
    reveal_type(d[0])  # revealed: Literal[0]
    reveal_type(b[0])  # revealed: Literal[0]
    reveal_type(dd[0])  # revealed: Literal[0]
    # TODO: should be Literal[0]
    reveal_type(cm[0])  # revealed: Unknown

[reveal_type(l[0]) for _ in range(1)]  # revealed: Literal[0]
[reveal_type(d[0]) for _ in range(1)]  # revealed: Literal[0]
[reveal_type(b[0]) for _ in range(1)]  # revealed: Literal[0]
[reveal_type(dd[0]) for _ in range(1)]  # revealed: Literal[0]
# TODO: should be Literal[0]
[reveal_type(cm[0]) for _ in range(1)]  # revealed: Unknown

def _():
    reveal_type(l[0])  # revealed: int | None
    reveal_type(d[0])  # revealed: int
    reveal_type(b[0])  # revealed: int
    reveal_type(dd[0])  # revealed: int
    reveal_type(cm[0])  # revealed: int

class D(TypedDict):
    x: int
    label: str

td = D(x=1, label="a")
td["x"] = 0
# TODO: should be Literal[0]
reveal_type(td["x"])  # revealed: @Todo(TypedDict)

# error: [unresolved-reference]
does["not"]["exist"] = 0
# error: [unresolved-reference]
reveal_type(does["not"]["exist"])  # revealed: Unknown

non_subscriptable = 1
# error: [non-subscriptable]
non_subscriptable[0] = 0
# error: [non-subscriptable]
reveal_type(non_subscriptable[0])  # revealed: Unknown

No narrowing for custom classes with arbitrary __getitem__ / __setitem__

class C:
    def __init__(self):
        self.l: list[str] = []

    def __getitem__(self, index: int) -> str:
        return self.l[index]

    def __setitem__(self, index: int, value: str | int) -> None:
        if len(self.l) == index:
            self.l.append(str(value))
        else:
            self.l[index] = str(value)

c = C()
c[0] = 0
reveal_type(c[0])  # revealed: str

Complex target

class A:
    x: list[int | None] = []

class B:
    a: A | None = None

b = B()
b.a = A()
b.a.x[0] = 0

reveal_type(b.a.x[0])  # revealed: Literal[0]

class C:
    reveal_type(b.a.x[0])  # revealed: Literal[0]

def _():
    # error: [possibly-unbound-attribute]
    reveal_type(b.a.x[0])  # revealed: Unknown | int | None
    # error: [possibly-unbound-attribute]
    reveal_type(b.a.x)  # revealed: Unknown | list[int | None]
    reveal_type(b.a)  # revealed: Unknown | A | None

Invalid assignments are not used for narrowing

class C:
    x: int | None
    l: list[int]

def f(c: C, s: str):
    c.x = s  # error: [invalid-assignment]
    reveal_type(c.x)  # revealed: int | None
    s = c.x  # error: [invalid-assignment]

    # TODO: This assignment is invalid and should result in an error.
    c.l[0] = s
    reveal_type(c.l[0])  # revealed: int