diff --git a/crates/ty_python_semantic/resources/mdtest/assignment/augmented.md b/crates/ty_python_semantic/resources/mdtest/assignment/augmented.md index 1787834ee654e3..9c368d9618f72a 100644 --- a/crates/ty_python_semantic/resources/mdtest/assignment/augmented.md +++ b/crates/ty_python_semantic/resources/mdtest/assignment/augmented.md @@ -187,6 +187,371 @@ def f(flag: bool, flag2: bool): reveal_type(f) # revealed: float | str ``` +## Declared attributes with in-place operators + +`+=` assigns the value returned by `__iadd__` back to its target. That value must be compatible with +the attribute's declared type. + +```py +class Value: + def __iadd__(self, other: int) -> str: + return "updated" + +class Holder: + value: Value + +holder = Holder() +# error: [invalid-assignment] +holder.value += 1 +reveal_type(holder.value) # revealed: Value +``` + +## Declared attributes without in-place operators + +When an object does not define `__iadd__`, `+=` falls back to `__add__`. Its result must still be +compatible with the attribute's declared type. + +```py +class Value: + def __add__(self, other: int) -> str: + return "updated" + +class Holder: + value: Value + +holder = Holder() +# error: [invalid-assignment] +holder.value += 1 +``` + +## Inferred attributes in loops + +An unannotated instance attribute may change type. After its initial `None` value is replaced, an +augmented assignment inside a loop must also contribute its result to the inferred attribute type. + +```py +class Counter: + def update(self) -> None: + self.value = None + self.value = 0 + for _ in range(1): + self.value += 1.0 + +reveal_type(Counter().value) # revealed: None | float +``` + +## Inferred class attributes + +An unannotated class attribute still has an inferred type that restricts assignments through an +instance. + +```py +class Holder: + value = 1 + +holder = Holder() +# error: [invalid-assignment] +holder.value += 0.5 +``` + +## Read-only properties + +`+=` writes its result back to the attribute. A property without a setter therefore cannot be the +target of an augmented assignment. + +```py +class ReadOnly: + @property + def value(self) -> int: + return 1 + +read_only = ReadOnly() +# error: [invalid-assignment] +read_only.value += 1 +``` + +## Properties with different getter and setter types + +A property can accept a wider type in its setter than it returns from its getter. The result of `/=` +is checked against the setter, while subsequent reads still use the getter's return type. + +```py +class Counter: + @property + def value(self) -> int: + return 1 + + @value.setter + def value(self, value: float) -> None: + pass + +counter = Counter() +counter.value /= 2 +reveal_type(counter.value) # revealed: int +``` + +## Attributes defined by descriptors + +When an unannotated class attribute is a data descriptor, its `__set__` method determines which +values may be assigned. + +```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 Holder: + value = Descriptor() + +holder = Holder() +# error: [invalid-assignment] +holder.value += 1 +``` + +## Custom subscript assignments + +`/=` first reads an item, then writes the result back through `__setitem__`. The assigned value is +the result of the operation, not the right-hand operand. + +```py +class Container: + def __getitem__(self, key: int) -> int: + return 1 + + def __setitem__(self, key: int, value: int) -> None: + pass + +container = Container() +# error: [invalid-assignment] +container[0] /= 2 +reveal_type(container[0]) # revealed: int +``` + +## Subscript setters with different value types + +A collection can accept a wider type in `__setitem__` than `__getitem__` returns. After a valid +assignment, subsequent reads still use the return type of `__getitem__`. + +```py +class Container: + def __getitem__(self, key: int) -> int: + return 1 + + def __setitem__(self, key: int, value: float) -> None: + pass + +container = Container() +container[0] /= 2 +reveal_type(container[0]) # revealed: int +``` + +## Annotated collection entries + +An annotation fixes the element type of a list, so `/=` cannot write a `float` into a `list[int]`. + +```py +values: list[int] = [1] +# error: [invalid-assignment] +values[0] /= 2 +``` + +The same rule applies to the value type of an annotated dictionary. + +```py +mapping: dict[str, int] = {"value": 1} +# error: [invalid-assignment] +mapping["value"] /= 2 +``` + +An annotated collection remains constrained when it is accessed through an attribute. + +```py +class Holder: + values: list[int] + +holder = Holder() +# error: [invalid-assignment] +holder.values[0] /= 2 +``` + +## Typed dictionary entries + +A `TypedDict` field can only be assigned a value compatible with its declared type. + +```py +from typing import TypedDict + +class Payload(TypedDict): + value: int + +payload: Payload = {"value": 1} +# error: [invalid-assignment] +payload["value"] /= 2 +``` + +## Read-only subscripts + +A readable item cannot be reassigned when its container does not implement `__setitem__`. + +```py +values: tuple[int] = (1,) +# error: [invalid-assignment] +values[0] += 1 +``` + +## Missing attributes + +If an augmented assignment cannot read its target, it must report that failure only once; no +assignment is attempted. + +```py +class Missing: ... + +missing = Missing() +# error: [unresolved-attribute] +missing.value += 1 +``` + +The same applies when an attribute is missing from one member of a union. + +```py +class Counter: + count: int + +def update(counter: Counter | None) -> None: + # error: [unresolved-attribute] + counter.count += 1 +``` + +## Invalid subscript reads + +An invalid key prevents an item from being read, so the failed assignment must not produce a second +error. + +```py +mapping: dict[str, int] = {} +# error: [invalid-argument-type] +mapping[1] += 1 +``` + +A value without `__getitem__` also fails before assignment can be attempted. + +```py +value = 1 +# error: [not-subscriptable] +value[0] += 1 +``` + +## Right-hand-side errors after failed reads + +Even when an attribute cannot be read, the right-hand side must still be checked for unrelated +errors. + +```py +class Missing: ... + +missing = Missing() +# error: [unresolved-attribute] +# error: [unresolved-reference] +missing.value += missing_attribute_operand +``` + +The same rule applies when a subscript cannot be read. + +```py +mapping: dict[str, int] = {} +# error: [invalid-argument-type] +# error: [unresolved-reference] +mapping[1] += missing_subscript_operand +``` + +## Failed in-place operations + +If `__iadd__` rejects its operand, its return type must not be treated as a value to assign. + +```py +class Value: + def __iadd__(self, other: int) -> str: + return "updated" + +class Holder: + value: Value + +holder = Holder() +# error: [unsupported-operator] +holder.value += "invalid" +``` + +## Union attribute assignments + +When objects in a union have different attribute types, each operator result should be checked +against the attribute from the same object. Ordinary assignments already lose this relationship, so +augmented assignments currently report the same false positive. + +```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: Check each result against the attribute it came from. + # error: [invalid-assignment] + value.value += 1 +``` + +## Collections that may be read-only + +When a collection could be a writable list or a read-only tuple, an item assignment is invalid +because it cannot be performed on every possible value. + +```py +def update(value: list[int] | tuple[int, ...]) -> None: + # error: [invalid-assignment] + value[0] += 1 +``` + +## Typed dictionary assignments with multiple possible keys + +A key that can select fields with different value types must only be assigned a value accepted by +every possible field. + +```py +from typing import Literal, TypedDict + +class Payload(TypedDict): + whole: int + fractional: float + +def update(value: Payload, key: Literal["whole", "fractional"]) -> None: + # error: [invalid-assignment] + value[key] /= 2 +``` + +## Inferred collection entries + +Augmented assignments are not yet included when inferring the element type of an unannotated +collection. + +```py +values = [1] +# TODO: Infer `list[float]` instead of rejecting the assignment. +# error: [invalid-assignment] +values[0] /= 2 +``` + ## Implicit dunder calls on class objects ```py diff --git a/crates/ty_python_semantic/resources/mdtest/attributes.md b/crates/ty_python_semantic/resources/mdtest/attributes.md index b488a2f05fb982..016f1f2c277cd5 100644 --- a/crates/ty_python_semantic/resources/mdtest/attributes.md +++ b/crates/ty_python_semantic/resources/mdtest/attributes.md @@ -259,6 +259,9 @@ reveal_type(c_instance.b) # revealed: int #### Augmented assignments +An augmented assignment contributes its result to the inferred type of an unannotated instance +attribute. + ```py class Weird: def __iadd__(self, other: None) -> str: @@ -269,9 +272,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 only `str`, since the initial `Weird` value has been overwritten. +reveal_type(C().w) # revealed: Weird | str ``` #### Nested augmented assignments after narrowing diff --git a/crates/ty_python_semantic/src/types/class/static_literal.rs b/crates/ty_python_semantic/src/types/class/static_literal.rs index e36a3935af95c1..b94428ddf531b0 100644 --- a/crates/ty_python_semantic/src/types/class/static_literal.rs +++ b/crates/ty_python_semantic/src/types/class/static_literal.rs @@ -48,7 +48,7 @@ use crate::{ is_implicit_staticmethod, }, generics::Specialization, - infer::infer_unpack_types, + infer::{infer_definition_types, infer_unpack_types}, infer_expression_type, inferred_declaration, known_instance::DeprecatedInstance, member::{Member, class_member}, @@ -3079,8 +3079,7 @@ impl<'db> StaticClassLiteral<'db> { } } DefinitionKind::AugmentedAssignment(_) => { - // TODO: - None + Some(infer_definition_types(db, binding).binding_type(binding)) } DefinitionKind::NamedExpression(_) => { // A named expression whose target is an attribute is syntactically prohibited diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index b8c5448e318ce2..9b3463a74481ed 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -1485,7 +1485,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { fn fallback_member_declared_type(&mut self, node: AnyNodeRef<'_>) -> Option> { let db = self.db(); if let AnyNodeRef::ExprAttribute(ast::ExprAttribute { value, attr, .. }) = node { - let value_type = self.infer_maybe_standalone_expression(value, TypeContext::default()); + let value_type = self.try_expression_type(value).unwrap_or_else(|| { + self.infer_maybe_standalone_expression(value, TypeContext::default()) + }); if let Place::Defined(DefinedPlace { ty, definedness: Definedness::AlwaysDefined, @@ -1505,9 +1507,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { }, ) = node { - let value_ty = self.infer_expression(value, TypeContext::default()); - let slice_ty = self.infer_expression(slice, TypeContext::default()); - Some(self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx)) + let value_ty = self.get_or_infer_expression(value, TypeContext::default()); + let slice_ty = self.get_or_infer_expression(slice, TypeContext::default()); + Some( + self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx) + .unwrap_or_else(|recovery_ty| recovery_ty), + ) } else { None } @@ -3213,13 +3218,24 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } ast::Expr::Subscript(subscript_expr) => { if let Some(infer_assigned_ty) = infer_assigned_ty { + let object_ty = + self.infer_expression(&subscript_expr.value, TypeContext::default()); + let mut infer_slice_ty = |builder: &mut Self, tcx| { + builder.infer_expression(&subscript_expr.slice, tcx) + }; let infer_assigned_ty = &mut |builder: &mut Self, tcx| { let assigned_ty = infer_assigned_ty(builder, tcx); builder.store_expression_type(target, assigned_ty); assigned_ty }; - self.validate_subscript_assignment(subscript_expr, value, infer_assigned_ty); + self.validate_subscript_assignment( + subscript_expr, + value, + object_ty, + &mut infer_slice_ty, + infer_assigned_ty, + ); } } @@ -4630,23 +4646,49 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { self.infer_definition(assignment); } else { // Non-name assignment targets are inferred as ordinary expressions, not definitions. - self.infer_augment_assignment(assignment); + if let Ok(result_ty) = self.infer_augment_assignment(assignment) { + let target = assignment.target.as_ref(); + match target { + ast::Expr::Attribute(attribute) => { + let object_ty = self.expression_type(&attribute.value); + self.validate_attribute_assignment( + attribute, + target, + object_ty, + attribute.attr.id(), + &mut |_, _| result_ty, + true, + ); + } + ast::Expr::Subscript(subscript) => { + let object_ty = self.expression_type(&subscript.value); + let slice_ty = self.expression_type(&subscript.slice); + self.validate_subscript_assignment( + subscript, + target, + object_ty, + &mut |_, _| slice_ty, + &mut |_, _| result_ty, + ); + } + _ => {} + } + } if let ast::Expr::Attribute(attr_expr) = assignment.target.as_ref() { - let object_ty = self.expression_type(&attr_expr.value); self.report_undeclared_protocol_attribute(attr_expr); - self.validate_final_attribute_assignment(attr_expr, object_ty, attr_expr.attr.id()); } } } + /// Infer an augmented operator, returning its recovery type if the operation fails. fn infer_augmented_op( &mut self, assignment: &ast::StmtAugAssign, target_type: Type<'db>, value_expr: &ast::Expr, infer_value_ty: &mut dyn FnMut(&mut Self, TypeContext<'db>) -> Type<'db>, - ) -> Type<'db> { + ) -> Result, Type<'db>> { let db = self.db(); let env = self.program_environment(); // If the target defines, e.g., `__iadd__`, infer the augmented assignment as a call to that @@ -4657,7 +4699,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let binary_return_ty = |builder: &mut Self, value_ty| { builder .infer_binary_expression_type(assignment.into(), false, target_type, value_ty, op) - .unwrap_or_else(|| { + .ok_or_else(|| { report_unsupported_augmented_assignment( &builder.context, assignment, @@ -4676,14 +4718,27 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // equally applicable type contexts for each union member. infer_value_ty.infer_loud(self, TypeContext::default()); - union.map(db, env, |&elem_type| { - self.infer_augmented_op( + let mut operation_failed = false; + let result_ty = union.map(db, env, |&elem_type| { + match self.infer_augmented_op( assignment, elem_type, value_expr, &mut |builder, tcx| infer_value_ty.infer_silent(builder, tcx), - ) - }) + ) { + Ok(ty) => ty, + Err(recovery_ty) => { + operation_failed = true; + recovery_ty + } + } + }); + + if operation_failed { + Err(result_ty) + } else { + Ok(result_ty) + } } _ => { @@ -4695,7 +4750,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { infer_value_ty, ) { - return typed_dict_update_ty; + return Ok(typed_dict_update_ty); } let ast_arguments = [ArgOrKeyword::Arg(value_expr)]; @@ -4711,7 +4766,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { TypeContext::default(), ); match call { - Ok(outcome) => outcome.return_type(db, env), + Ok(outcome) => Ok(outcome.return_type(db, env)), Err(CallDunderError::MethodNotAvailable) => { let value_ty = infer_value_ty(self, TypeContext::default()); binary_return_ty(self, value_ty) @@ -4720,12 +4775,20 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { bindings: outcome, .. }) => { let value_ty = outcome.type_for_argument(&call_arguments, 0); - UnionType::from_two_elements( - db, - env, - outcome.return_type(db, env), - binary_return_ty(self, value_ty), - ) + match binary_return_ty(self, value_ty) { + Ok(binary_ty) => Ok(UnionType::from_two_elements( + db, + env, + outcome.return_type(db, env), + binary_ty, + )), + Err(recovery_ty) => Err(UnionType::from_two_elements( + db, + env, + outcome.return_type(db, env), + recovery_ty, + )), + } } Err(CallDunderError::CallError(_, bindings, _)) => { let value_ty = bindings.type_for_argument(&call_arguments, 0); @@ -4735,7 +4798,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { target_type, value_ty, ); - bindings.return_type(db, env) + Err(bindings.return_type(db, env)) } } } @@ -4747,12 +4810,17 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { assignment: &'ast ast::StmtAugAssign, definition: Definition<'db>, ) { - let target_ty = self.infer_augment_assignment(assignment); - self.add_binding(assignment.into(), definition) + let target_ty = self + .infer_augment_assignment(assignment) + .unwrap_or_else(|recovery_ty| recovery_ty); + self.add_binding(assignment.target.as_ref().into(), definition) .insert(self, target_ty); } - fn infer_augment_assignment(&mut self, assignment: &ast::StmtAugAssign) -> Type<'db> { + fn infer_augment_assignment( + &mut self, + assignment: &ast::StmtAugAssign, + ) -> Result, Type<'db>> { let ast::StmtAugAssign { range: _, node_index: _, @@ -4762,28 +4830,37 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } = assignment; // Resolve the target type, assuming a load context. - let target_type = match &**target { + let target_result = match &**target { ast::Expr::Name(name) => { let previous_value = self.infer_name_load(name); self.store_expression_type(target, previous_value); - previous_value + Ok(previous_value) } ast::Expr::Attribute(attr) => { - let previous_value = self.infer_attribute_load(attr); + let result = self.infer_attribute_load(attr); + let previous_value = result.unwrap_or_else(|recovery_ty| recovery_ty); self.store_expression_type(target, previous_value); - previous_value + result } ast::Expr::Subscript(subscript) => { - let previous_value = self.infer_subscript_load(subscript); + let result = self.infer_subscript_load(subscript); + let previous_value = result.unwrap_or_else(|recovery_ty| recovery_ty); self.store_expression_type(target, previous_value); - previous_value + result } - _ => self.infer_expression(target, TypeContext::default()), + _ => Ok(self.infer_expression(target, TypeContext::default())), }; - self.infer_augmented_op(assignment, target_type, value, &mut |builder, tcx| { - builder.infer_expression(value, tcx) - }) + let target_type = target_result.unwrap_or_else(|recovery_ty| recovery_ty); + let operation_result = + self.infer_augmented_op(assignment, target_type, value, &mut |builder, tcx| { + builder.infer_expression(value, tcx) + }); + + match (target_result, operation_result) { + (Ok(_), Ok(result_ty)) => Ok(result_ty), + (_, Ok(recovery_ty) | Err(recovery_ty)) => Err(recovery_ty), + } } fn infer_dict_key_assignment_definition( @@ -9193,6 +9270,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let collection_generic_context = collection_literal.generic_context(db); let mut identity_bindings = self .infer_attribute_load_impl(attribute, identity_instance) + .unwrap_or_else(|recovery_ty| recovery_ty) .bindings(db, env) .match_parameters(db, env, &call_arguments) // Perform inference against the type variables on the receiver's generic context. @@ -10253,19 +10331,22 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - /// Infer the type of a [`ast::ExprAttribute`] expression, assuming a load context. - fn infer_attribute_load(&mut self, attribute: &ast::ExprAttribute) -> Type<'db> { + /// Infer an attribute load, returning its recovery type if lookup fails. + fn infer_attribute_load( + &mut self, + attribute: &ast::ExprAttribute, + ) -> Result, Type<'db>> { let value_type = self.infer_maybe_standalone_expression(&attribute.value, TypeContext::default()); self.infer_attribute_load_impl(attribute, value_type) } - /// Infer the type of a [`ast::ExprAttribute`] expression, assuming a load context. + /// Infer an attribute load on a known receiver, returning its recovery type if lookup fails. fn infer_attribute_load_impl( &mut self, attribute: &ast::ExprAttribute, mut value_type: Type<'db>, - ) -> Type<'db> { + ) -> Result, Type<'db>> { fn union_elements_missing_attribute<'db>( db: &'db dyn Db, env: &ProgramEnvironment<'db>, @@ -10328,8 +10409,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { }); let attr_name = &attr.id; - let resolved_type = - fallback_place.unwrap_with_diagnostic(db, env, |lookup_err| match lookup_err { + let lookup_result = fallback_place.into_lookup_result(db, env); + let resolved_type = lookup_result.unwrap_or_else(|lookup_err| match lookup_err { LookupError::Undefined(_) => { let fallback = || { TypeAndQualifiers::new( @@ -10586,7 +10667,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { type_when_bound } - }); + }); let resolved_type = resolved_type.inner_type(); @@ -10594,7 +10675,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // Even if we can obtain the attribute type based on the assignments, we still perform default type inference // (to report errors). - assigned_type.unwrap_or(resolved_type) + let inferred_type = assigned_type.unwrap_or(resolved_type); + lookup_result + .map(|_| inferred_type) + .map_err(|_| inferred_type) } fn infer_attribute_expression(&mut self, attribute: &ast::ExprAttribute) -> Type<'db> { @@ -10607,13 +10691,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } = attribute; match ctx { - ExprContext::Load => self.infer_attribute_load(attribute), + ExprContext::Load => self + .infer_attribute_load(attribute) + .unwrap_or_else(|recovery_ty| recovery_ty), ExprContext::Store => { self.infer_expression(value, TypeContext::default()); Type::Never } ExprContext::Del => { - self.infer_attribute_load(attribute); + let _ = self.infer_attribute_load(attribute); self.validate_attribute_deletion( attribute, self.expression_type(value), diff --git a/crates/ty_python_semantic/src/types/infer/builder/subscript.rs b/crates/ty_python_semantic/src/types/infer/builder/subscript.rs index d5f9680793d18d..f5f9b99e2f4151 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/subscript.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/subscript.rs @@ -136,12 +136,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } = subscript; match ctx { - ExprContext::Load => self.infer_subscript_load(subscript), + ExprContext::Load => self + .infer_subscript_load(subscript) + .unwrap_or_else(|recovery_ty| recovery_ty), ExprContext::Store => { let value_ty = self.infer_expression(value, TypeContext::default()); self.store_typed_dict_key_expected_type(slice, value_ty); let slice_ty = self.infer_expression(slice, TypeContext::default()); - self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); + let _ = self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); Type::Never } ExprContext::Del => { @@ -154,20 +156,24 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { ExprContext::Invalid => { let value_ty = self.infer_expression(value, TypeContext::default()); let slice_ty = self.infer_expression(slice, TypeContext::default()); - self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); + let _ = self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); Type::unknown() } } } - pub(super) fn infer_subscript_load(&mut self, subscript: &ast::ExprSubscript) -> Type<'db> { + /// Infer a subscript load, returning its recovery type if the subscription fails. + pub(super) fn infer_subscript_load( + &mut self, + subscript: &ast::ExprSubscript, + ) -> Result, Type<'db>> { let value_ty = self.infer_expression(&subscript.value, TypeContext::default()); // If we have an implicit type alias like `MyList = list[T]`, and if `MyList` is being // used in another implicit type alias like `Numbers = MyList[int]`, then we infer the // right hand side as a value expression, and need to handle the specialization here. if value_ty.is_generic_alias() { - return self.infer_explicit_type_alias_specialization(subscript, value_ty, false); + return Ok(self.infer_explicit_type_alias_specialization(subscript, value_ty, false)); } self.infer_subscript_load_impl(value_ty, subscript) @@ -177,7 +183,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { &mut self, value_ty: Type<'db>, subscript: &ast::ExprSubscript, - ) -> Type<'db> { + ) -> Result, Type<'db>> { let env = self.program_environment(); let db = self.db(); @@ -186,7 +192,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { node_index: _, value: _, slice, - ctx: expr_context, + ctx: _, } = subscript; self.store_typed_dict_key_expected_type(slice, value_ty); @@ -210,13 +216,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // Even if we can obtain the subscript type based on the assignments, we still perform default type inference // (to store the expression type and to report errors). let slice_ty = self.infer_expression(slice, TypeContext::default()); - self.infer_subscript_expression_types( - subscript, - value_ty, - slice_ty, - *expr_context, - ); - return ty; + return self + .infer_subscript_expression_types( + subscript, + value_ty, + slice_ty, + ExprContext::Load, + ) + .map(|_| ty) + .map_err(|_| ty); } } } @@ -235,43 +243,49 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // updating all of the subscript logic below to use custom callables for all of the _other_ // special cases, too. if class.is_tuple(db) { - return tuple_generic_alias(env, self.infer_tuple_type_expression(subscript)); + return Ok(tuple_generic_alias( + env, + self.infer_tuple_type_expression(subscript), + )); } else if class.is_known(db, KnownClass::Type) { let argument_ty = self.infer_type_expression(slice); - return Type::KnownInstance(KnownInstanceType::TypeGenericAlias( + return Ok(Type::KnownInstance(KnownInstanceType::TypeGenericAlias( InternedType::new(db, argument_ty), - )); + ))); } if let Some(generic_context) = class.generic_context(db) && let Some(class) = class.as_static() { - return self.infer_explicit_class_specialization( + return Ok(self.infer_explicit_class_specialization( subscript, value_ty, class, generic_context, - ); + )); } } Type::KnownInstance(KnownInstanceType::TypeAliasType(type_alias)) => { if let Some(generic_context) = type_alias.generic_context(db) { - return self.infer_explicit_type_alias_type_specialization( + return Ok(self.infer_explicit_type_alias_type_specialization( subscript, value_ty, type_alias, generic_context, - ); + )); } } Type::SpecialForm(special_form) => match special_form { SpecialFormType::Tuple => { - return tuple_generic_alias(env, self.infer_tuple_type_expression(subscript)); + return Ok(tuple_generic_alias( + env, + self.infer_tuple_type_expression(subscript), + )); } SpecialFormType::Literal => match self.infer_literal_parameter_type(slice) { Ok(result) => { - return Type::KnownInstance(KnownInstanceType::Literal(InternedType::new( - db, result, + return Ok(Type::KnownInstance(KnownInstanceType::Literal( + InternedType::new(db, result), ))); } Err(nodes) => { @@ -285,16 +299,16 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { a literal value (int, bool, str, or bytes), or an enum member", ); } - return Type::unknown(); + return Ok(Type::unknown()); } }, SpecialFormType::Annotated => { - return self + return Ok(self .parse_subscription_of_annotated_special_form( subscript, AnnotatedExprContext::TypeExpression, ) - .inner_type(); + .inner_type()); } SpecialFormType::Optional => { if matches!(**slice, ast::Expr::Tuple(_)) @@ -310,9 +324,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // `Optional[None]` is equivalent to `None`: if ty.is_none(db) { - return ty; + return Ok(ty); } - return Type::KnownInstance(KnownInstanceType::UnionType( + return Ok(Type::KnownInstance(KnownInstanceType::UnionType( UnionTypeInstance::new( db, None, @@ -323,7 +337,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { Type::none(db, env), )), ), - )); + ))); } SpecialFormType::Union => match **slice { ast::Expr::Tuple(ref tuple) => { @@ -346,18 +360,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { ); } - return union_type; + return Ok(union_type); } _ => { - return self.infer_expression(slice, TypeContext::default()); + return Ok(self.infer_expression(slice, TypeContext::default())); } }, SpecialFormType::Type => { // Similar to the branch above that handles `type[…]`, handle `typing.Type[…]` let argument_ty = self.infer_type_expression(slice); - return Type::KnownInstance(KnownInstanceType::TypeGenericAlias( + return Ok(Type::KnownInstance(KnownInstanceType::TypeGenericAlias( InternedType::new(db, argument_ty), - )); + ))); } SpecialFormType::TypingCallable | SpecialFormType::CollectionsAbcCallable => { let callable = self @@ -365,7 +379,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .as_callable() .expect("always returns Type::Callable"); - return Type::KnownInstance(KnownInstanceType::Callable(callable)); + return Ok(Type::KnownInstance(KnownInstanceType::Callable(callable))); } SpecialFormType::Unpack => { self.store_type_expression_flags( @@ -383,19 +397,21 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { previously_in_unpack_type_argument, ); - return if matches!( - inner_ty, - Type::TypeVar(typevar) if typevar.is_typevartuple(db) - ) || inner_ty.exact_tuple_instance_spec(db).is_some() - { - inner_ty - } else { - self.store_type_expression_flags( - ast::ExprRef::from(subscript), - TypeExpressionFlags::INVALID_UNPACK, - ); - Type::unknown() - }; + return Ok( + if matches!( + inner_ty, + Type::TypeVar(typevar) if typevar.is_typevartuple(db) + ) || inner_ty.exact_tuple_instance_spec(db).is_some() + { + inner_ty + } else { + self.store_type_expression_flags( + ast::ExprRef::from(subscript), + TypeExpressionFlags::INVALID_UNPACK, + ); + Type::unknown() + }, + ); } SpecialFormType::LegacyStdlibAlias(alias) => { let AliasSpec { @@ -431,10 +447,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .map(|arg| self.infer_type_expression(arg)) .collect(); - return class + return Ok(class .to_specialized_class_type(db, env, arg_types) .map(Type::from) - .unwrap_or_else(Type::unknown); + .unwrap_or_else(Type::unknown)); } _ => {} }, @@ -445,7 +461,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { | KnownInstanceType::Callable(_) | KnownInstanceType::TypeGenericAlias(_), ) => { - return self.infer_explicit_type_alias_specialization(subscript, value_ty, false); + return Ok( + self.infer_explicit_type_alias_specialization(subscript, value_ty, false) + ); } Type::Dynamic(DynamicType::Unknown) => { let slice_ty = self.infer_expression(slice, TypeContext::default()); @@ -457,15 +475,21 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { &mut variables, ); let generic_context = GenericContext::from_typevar_instances(db, env, variables); - return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); + return Ok(Type::Dynamic(DynamicType::UnknownGeneric(generic_context))); } _ => {} } let slice_ty = self.infer_expression(slice, TypeContext::default()); - let result_ty = - self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *expr_context); - self.narrow_expr_with_applicable_constraints(subscript, result_ty, &constraint_keys) + self.infer_subscript_expression_types(subscript, value_ty, slice_ty, ExprContext::Load) + .map(|ty| self.narrow_expr_with_applicable_constraints(subscript, ty, &constraint_keys)) + .map_err(|recovery_ty| { + self.narrow_expr_with_applicable_constraints( + subscript, + recovery_ty, + &constraint_keys, + ) + }) } pub(super) fn infer_explicit_class_specialization( @@ -1396,13 +1420,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { Err(()) } + /// Infer a subscription and report failures while preserving their recovery types. pub(super) fn infer_subscript_expression_types( &self, subscript: &ast::ExprSubscript, value_ty: Type<'db>, slice_ty: Type<'db>, expr_context: ExprContext, - ) -> Type<'db> { + ) -> Result, Type<'db>> { let env = self.program_environment(); let db = self.db(); @@ -1462,7 +1487,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { SubscriptErrorKind::MultipleTypeVarTuples { origin }, ); error.report_diagnostics(&self.context, subscript); - return error.result_type(); + return Err(error.result_type()); } if has_invalid_unpack_argument { let error = SubscriptError::new( @@ -1473,7 +1498,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { }, ); error.report_diagnostics(&self.context, subscript); - return error.result_type(); + return Err(error.result_type()); } } @@ -1516,9 +1541,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { _ => value_ty.subscript(db, env, slice_ty, expr_context), }; - subscript_result.unwrap_or_else(|e| { - e.report_diagnostics(&self.context, subscript); - e.result_type() + subscript_result.map_err(|error| { + error.report_diagnostics(&self.context, subscript); + error.result_type() }) } @@ -1553,6 +1578,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { &mut self, target: &ast::ExprSubscript, rhs_value: &ast::Expr, + object_ty: Type<'db>, + infer_slice_ty: &mut dyn FnMut(&mut Self, TypeContext<'db>) -> Type<'db>, infer_rhs_value: &mut dyn FnMut(&mut Self, TypeContext<'db>) -> Type<'db>, ) -> bool { let env = self.program_environment(); @@ -1566,15 +1593,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let db = self.db(); - let object_ty = self.infer_expression(object, TypeContext::default()); self.store_typed_dict_key_expected_type(slice, object_ty); - let mut infer_slice_ty = |builder: &mut Self, tcx| builder.infer_expression(slice, tcx); let is_valid_assignment = self.validate_subscript_assignment_impl( target, None, object_ty, - &mut infer_slice_ty, + infer_slice_ty, rhs_value, infer_rhs_value, true,