[ty] Infer collection types from augmented subscript stores

crates/ty_python_core/src/builder.rs

@@ -4528,6 +4528,17 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                         }
                     }
 
+                    // An augmented single-element assignment on the collection object. Slices
+                    // produce another collection and require a separate constraint model.
+                    ruff_python_ast::Stmt::AugAssign(ast::StmtAugAssign { target, .. }) => {
+                        matches!(
+                            target.as_ref(),
+                            ast::Expr::Subscript(ast::ExprSubscript { value, slice, .. })
+                                if !matches!(slice.as_ref(), ast::Expr::Slice(_))
+                                    && ExpressionNodeKey::from(value) == *use_expression
+                        )
+                    }
+
                     // An annotated assignment assigning the collection object to a new binding.
                     ruff_python_ast::Stmt::AnnAssign(_) => true,
 

crates/ty_python_semantic/resources/mdtest/assignment/augmented.md

@@ -187,6 +187,329 @@ def f(flag: bool, flag2: bool):
     reveal_type(f)  # revealed: float | str
 ```
 
+## Annotated name targets
+
+An augmented assignment to an annotated name must validate its result against the declaration. An
+unannotated name can instead change type.
+
+```py
+class Value:
+    def __add__(self, other: int) -> object:
+        return other
+
+annotated: Value = Value()
+# error: [invalid-assignment]
+annotated += 1
+reveal_type(annotated)  # revealed: Value
+
+inferred = Value()
+inferred += 1
+reveal_type(inferred)  # revealed: object
+```
+
+## Attribute targets
+
+The result must satisfy the attribute's write contract, whether the operation uses `__iadd__` or
+falls back to `__add__`.
+
+```py
+class AddValue:
+    def __add__(self, other: int) -> object:
+        return other
+
+class InplaceValue:
+    def __iadd__(self, other: int) -> object:
+        return other
+
+class Holder:
+    add: AddValue
+    inplace: InplaceValue
+
+holder = Holder()
+# error: [invalid-assignment]
+holder.add += 1
+reveal_type(holder.add)  # revealed: AddValue
+
+# error: [invalid-assignment]
+holder.inplace += 1
+reveal_type(holder.inplace)  # revealed: InplaceValue
+```
+
+## Inferred attribute targets in loops
+
+An inferred attribute can change type across assignments. Its initial value must not become a
+declaration that pollutes the loop-carried type after the attribute has been reassigned.
+
+```py
+class Counter:
+    def update(self, increment: float) -> None:
+        self.value = None
+        self.value = 0
+        for _ in range(1):
+            self.value += increment
+
+reveal_type(Counter().value)  # revealed: None | float
+```
+
+## Inferred public attribute targets
+
+An inferred class attribute has the same public write contract for augmented and ordinary
+assignments.
+
+```py
+class Holder:
+    value = 1
+
+holder = Holder()
+# error: [invalid-assignment]
+holder.value += 0.5
+```
+
+## Attribute descriptors
+
+Even an in-place operation writes its result back, so a read-only property rejects augmented
+assignment.
+
+```py
+class ReadOnly:
+    @property
+    def value(self) -> int:
+        return 1
+
+read_only = ReadOnly()
+# error: [invalid-assignment]
+read_only.value += 1
+reveal_type(read_only.value)  # revealed: int
+```
+
+A property's setter can accept a different type than its getter returns. The operation's result must
+satisfy the setter, while later reads continue to use the getter type.
+
+```py
+class ReadValue:
+    def __iadd__(self, other: int) -> str:
+        return "updated"
+
+class Writable:
+    @property
+    def value(self) -> ReadValue:
+        return ReadValue()
+
+    @value.setter
+    def value(self, value: str) -> None:
+        pass
+
+writable = Writable()
+writable.value += 1
+reveal_type(writable.value)  # revealed: ReadValue
+```
+
+Unannotated data descriptors still impose the write contract declared by their setter.
+
+```py
+class Descriptor:
+    def __get__(self, instance: object, owner: type[object] | None = None) -> int:
+        return 1
+
+    def __set__(self, instance: object, value: str) -> None:
+        pass
+
+class Custom:
+    value = Descriptor()
+
+custom = Custom()
+# error: [invalid-assignment]
+custom.value += 1
+```
+
+## Subscript targets
+
+An augmented subscript assignment must pass its result, not the operator's right-hand operand, to
+the target's `__setitem__` method.
+
+```py
+class Value:
+    def __iadd__(self, other: int) -> object:
+        return other
+
+class Container:
+    def __getitem__(self, key: int) -> Value:
+        return Value()
+
+    def __setitem__(self, key: int, value: Value) -> None:
+        pass
+
+container = Container()
+# error: [invalid-assignment]
+container[0] += 1
+reveal_type(container[0])  # revealed: Value
+```
+
+A custom setter may accept a broader type than its getter returns.
+
+```py
+class PermissiveContainer:
+    def __getitem__(self, key: int) -> Value:
+        return Value()
+
+    def __setitem__(self, key: int, value: object) -> None:
+        pass
+
+permissive = PermissiveContainer()
+permissive[0] += 1
+reveal_type(permissive[0])  # revealed: Value
+```
+
+Explicitly annotated lists and dictionaries retain their write contracts.
+
+```py
+items: list[Value] = [Value()]
+# error: [invalid-assignment]
+items[0] += 1
+reveal_type(items[0])  # revealed: Value
+
+mapping: dict[str, Value] = {"value": Value()}
+# error: [invalid-assignment]
+mapping["value"] += 1
+reveal_type(mapping["value"])  # revealed: Value
+```
+
+Declared collection-valued attributes also retain their write contracts.
+
+```py
+class Holder:
+    values: list[Value]
+
+holder = Holder()
+# error: [invalid-assignment]
+holder.values[0] += 1
+```
+
+Typed dictionary entries validate the value written back to their declared fields.
+
+```py
+from typing import TypedDict
+
+class Payload(TypedDict):
+    value: Value
+
+payload: Payload = {"value": Value()}
+# error: [invalid-assignment]
+payload["value"] += 1
+reveal_type(payload["value"])  # revealed: Value
+```
+
+## Read-only subscripts
+
+A readable subscript is not necessarily writable.
+
+```py
+values: tuple[int] = (1,)
+# error: [invalid-assignment]
+values[0] += 1
+```
+
+## Failed attribute and subscript loads
+
+Both the load and the store are checked, just as they are for an ordinary assignment whose value
+reads the same target.
+
+```py
+class Missing: ...
+
+missing = Missing()
+# error: [unresolved-attribute]
+# error: [unresolved-attribute]
+missing.value += 1
+
+mapping: dict[str, int] = {}
+# error: [invalid-argument-type]
+# error: [invalid-assignment]
+mapping[1] += 1
+```
+
+## Failed augmented operations
+
+An operation that cannot run does not perform a store.
+
+```py
+class Value:
+    def __iadd__(self, other: int) -> object:
+        return other
+
+class Holder:
+    value: Value
+
+holder = Holder()
+# error: [unsupported-operator]
+holder.value += "invalid"
+```
+
+## Correlated union targets
+
+The result of an operation on one union member must not be checked against another member's write
+contract.
+
+```py
+class AValue:
+    def __iadd__(self, other: int) -> "AValue":
+        return self
+
+class BValue:
+    def __iadd__(self, other: int) -> "BValue":
+        return self
+
+class A:
+    value: AValue
+
+class B:
+    value: BValue
+
+def update(value: A | B) -> None:
+    # TODO: Preserve receiver correlation, which is also lost in ordinary assignments.
+    # error: [invalid-assignment]
+    value.value += 1
+```
+
+## Union subscript targets
+
+An augmented assignment must reject a union alternative that does not support the write.
+
+```py
+def update(value: list[int] | tuple[int, ...]) -> None:
+    # error: [invalid-assignment]
+    value[0] += 1
+```
+
+## Union subscript keys
+
+Each possible typed-dictionary key must accept the value written by the augmented assignment.
+
+```py
+from typing import Literal, TypedDict
+
+class Payload(TypedDict):
+    first: int
+    second: int
+
+def update(value: Payload, key: Literal["first", "second"]) -> None:
+    value[key] += 1
+
+    # error: [invalid-assignment]
+    # error: [invalid-assignment]
+    value[key] /= 2
+```
+
+## Inferred collection targets
+
+Augmented subscript assignments contribute their operator result to full-scope collection inference.
+
+```py
+values = [1]
+values[0] /= 2
+reveal_type(values)  # revealed: list[float]
+```
+
 ## Implicit dunder calls on class objects
 
 ```py

crates/ty_python_semantic/resources/mdtest/attributes.md

@@ -269,9 +269,8 @@ class C:
         self.w = Weird()
         self.w += None
 
-# TODO: Mypy and pyright do not support this, but it would be great if we could
-# infer `str` here (`Weird` is not a possible type for the `w` attribute).
-reveal_type(C().w)  # revealed: Weird
+# TODO: Infer `str` alone, since the initial `Weird` value has been overwritten.
+reveal_type(C().w)  # revealed: Weird | str
 ```
 
 #### Nested augmented assignments after narrowing