diff --git a/crates/ty_python_semantic/resources/mdtest/assignment/annotations.md b/crates/ty_python_semantic/resources/mdtest/assignment/annotations.md index b5b328d219..8a0eb110e5 100644 --- a/crates/ty_python_semantic/resources/mdtest/assignment/annotations.md +++ b/crates/ty_python_semantic/resources/mdtest/assignment/annotations.md @@ -139,6 +139,15 @@ reveal_type(n) # revealed: list[Literal[1, 2, 3]] # error: [invalid-assignment] "Object of type `list[Unknown | str]` is not assignable to `list[LiteralString]`" o: list[typing.LiteralString] = ["a", "b", "c"] reveal_type(o) # revealed: list[LiteralString] + +p: dict[int, int] = {} +reveal_type(p) # revealed: dict[int, int] + +q: dict[int | str, int] = {1: 1, 2: 2, 3: 3} +reveal_type(q) # revealed: dict[int | str, int] + +r: dict[int | str, int | str] = {1: 1, 2: 2, 3: 3} +reveal_type(r) # revealed: dict[int | str, int | str] ``` ## Incorrect collection literal assignments are complained aobut diff --git a/crates/ty_python_semantic/resources/mdtest/call/builtins.md b/crates/ty_python_semantic/resources/mdtest/call/builtins.md index f32c007c1c..0eac021d1a 100644 --- a/crates/ty_python_semantic/resources/mdtest/call/builtins.md +++ b/crates/ty_python_semantic/resources/mdtest/call/builtins.md @@ -57,7 +57,7 @@ type("Foo", Base, {}) # error: [invalid-argument-type] "Argument to class `type` is incorrect: Expected `tuple[type, ...]`, found `tuple[Literal[1], Literal[2]]`" type("Foo", (1, 2), {}) -# TODO: this should be an error +# error: [invalid-argument-type] "Argument to class `type` is incorrect: Expected `dict[str, Any]`, found `dict[Unknown | bytes, Unknown | int]`" type("Foo", (Base,), {b"attr": 1}) ``` diff --git a/crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md b/crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md index 31154c77d0..7e1acf4efb 100644 --- a/crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md +++ b/crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md @@ -3,7 +3,49 @@ ## Empty dictionary ```py -reveal_type({}) # revealed: dict[@Todo(dict literal key type), @Todo(dict literal value type)] +reveal_type({}) # revealed: dict[Unknown, Unknown] +``` + +## Basic dict + +```py +reveal_type({1: 1, 2: 1}) # revealed: dict[Unknown | int, Unknown | int] +``` + +## Dict of tuples + +```py +reveal_type({1: (1, 2), 2: (3, 4)}) # revealed: dict[Unknown | int, Unknown | tuple[int, int]] +``` + +## Unpacked dict + +```py +a = {"a": 1, "b": 2} +b = {"c": 3, "d": 4} + +d = {**a, **b} +reveal_type(d) # revealed: dict[Unknown | str, Unknown | int] +``` + +## Dict of functions + +```py +def a(_: int) -> int: + return 0 + +def b(_: int) -> int: + return 1 + +x = {1: a, 2: b} +reveal_type(x) # revealed: dict[Unknown | int, Unknown | ((_: int) -> int)] +``` + +## Mixed dict + +```py +# revealed: dict[Unknown | str, Unknown | int | tuple[int, int] | tuple[int, int, int]] +reveal_type({"a": 1, "b": (1, 2), "c": (1, 2, 3)}) ``` ## Dict comprehensions diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md b/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md index bc2def2752..8f32950252 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md @@ -206,8 +206,7 @@ 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[@Todo(dict literal key type), @Todo(dict literal value type)] +reveal_type(cm) # revealed: ChainMap[Unknown | int, Unknown | int] reveal_type(l[0]) # revealed: Literal[0] reveal_type(d[0]) # revealed: Literal[0] diff --git a/crates/ty_python_semantic/resources/mdtest/typed_dict.md b/crates/ty_python_semantic/resources/mdtest/typed_dict.md index e44ed92633..3ad8df2eba 100644 --- a/crates/ty_python_semantic/resources/mdtest/typed_dict.md +++ b/crates/ty_python_semantic/resources/mdtest/typed_dict.md @@ -85,6 +85,34 @@ alice["extra"] = True bob["extra"] = True ``` +## Nested `TypedDict` + +Nested `TypedDict` fields are also supported. + +```py +from typing import TypedDict + +class Inner(TypedDict): + name: str + age: int | None + +class Person(TypedDict): + inner: Inner +``` + +```py +alice: Person = {"inner": {"name": "Alice", "age": 30}} + +reveal_type(alice["inner"]["name"]) # revealed: str +reveal_type(alice["inner"]["age"]) # revealed: int | None + +# error: [invalid-key] "Invalid key access on TypedDict `Inner`: Unknown key "non_existing"" +reveal_type(alice["inner"]["non_existing"]) # revealed: Unknown + +# error: [invalid-key] "Invalid key access on TypedDict `Inner`: Unknown key "extra"" +alice: Person = {"inner": {"name": "Alice", "age": 30, "extra": 1}} +``` + ## Validation of `TypedDict` construction ```py diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 2410d4ab32..b2cb2d61f1 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -849,6 +849,28 @@ impl<'db> Type<'db> { matches!(self, Type::Dynamic(_)) } + // If the type is a specialized instance of the given `KnownClass`, returns the specialization. + pub(crate) fn known_specialization( + self, + known_class: KnownClass, + db: &'db dyn Db, + ) -> Option> { + let class_type = match self { + Type::NominalInstance(instance) => instance, + Type::TypeAlias(alias) => alias.value_type(db).into_nominal_instance()?, + _ => return None, + } + .class(db); + + if !class_type.is_known(db, known_class) { + return None; + } + + class_type + .into_generic_alias() + .map(|generic_alias| generic_alias.specialization(db)) + } + /// Returns the top materialization (or upper bound materialization) of this type, which is the /// most general form of the type that is fully static. #[must_use] diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 647136b8db..22aeaba771 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -386,20 +386,8 @@ impl<'db> TypeContext<'db> { known_class: KnownClass, db: &'db dyn Db, ) -> Option> { - let class_type = match self.annotation? { - Type::NominalInstance(instance) => instance, - Type::TypeAlias(alias) => alias.value_type(db).into_nominal_instance()?, - _ => return None, - } - .class(db); - - if !class_type.is_known(db, known_class) { - return None; - } - - class_type - .into_generic_alias() - .map(|generic_alias| generic_alias.specialization(db)) + self.annotation + .and_then(|ty| ty.known_specialization(known_class, db)) } } diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index a276063415..6f20e78272 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -1,4 +1,6 @@ -use itertools::Itertools; +use std::iter; + +use itertools::{Either, Itertools}; use ruff_db::diagnostic::{Annotation, DiagnosticId, Severity}; use ruff_db::files::File; use ruff_db::parsed::ParsedModuleRef; @@ -86,13 +88,13 @@ use crate::types::typed_dict::{ }; use crate::types::visitor::any_over_type; use crate::types::{ - CallDunderError, CallableType, ClassLiteral, ClassType, DataclassParams, DynamicType, - IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType, MemberLookupPolicy, - MetaclassCandidate, PEP695TypeAliasType, Parameter, ParameterForm, Parameters, SpecialFormType, - SubclassOfType, TrackedConstraintSet, Truthiness, Type, TypeAliasType, TypeAndQualifiers, - TypeContext, TypeMapping, TypeQualifiers, TypeVarBoundOrConstraintsEvaluation, - TypeVarDefaultEvaluation, TypeVarInstance, TypeVarKind, UnionBuilder, UnionType, binding_type, - todo_type, + BoundTypeVarInstance, CallDunderError, CallableType, ClassLiteral, ClassType, DataclassParams, + DynamicType, IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType, + MemberLookupPolicy, MetaclassCandidate, PEP695TypeAliasType, Parameter, ParameterForm, + Parameters, SpecialFormType, SubclassOfType, TrackedConstraintSet, Truthiness, Type, + TypeAliasType, TypeAndQualifiers, TypeContext, TypeMapping, TypeQualifiers, + TypeVarBoundOrConstraintsEvaluation, TypeVarDefaultEvaluation, TypeVarInstance, TypeVarKind, + UnionBuilder, UnionType, binding_type, todo_type, }; use crate::types::{ClassBase, add_inferred_python_version_hint_to_diagnostic}; use crate::unpack::{EvaluationMode, UnpackPosition}; @@ -4110,7 +4112,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { value, TypeContext::new(Some(declared.inner_type())), ); - let mut inferred_ty = if target + let inferred_ty = if target .as_name_expr() .is_some_and(|name| &name.id == "TYPE_CHECKING") { @@ -4121,24 +4123,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { inferred_ty }; - // Validate `TypedDict` dictionary literal assignments - if let Some(typed_dict) = declared.inner_type().into_typed_dict() { - if let Some(dict_expr) = value.as_dict_expr() { - validate_typed_dict_dict_literal( - &self.context, - typed_dict, - dict_expr, - target.into(), - |expr| self.expression_type(expr), - ); - - // Override the inferred type of the dict literal to be the `TypedDict` type - // This ensures that the dict literal gets the correct type for key access - let typed_dict_type = Type::TypedDict(typed_dict); - inferred_ty = typed_dict_type; - } - } - self.add_declaration_with_binding( target.into(), definition, @@ -5290,6 +5274,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { ctx: _, } = list; + let elts = elts.iter().map(|elt| [Some(elt)]); self.infer_collection_literal(elts, tcx, KnownClass::List) .unwrap_or_else(|| { KnownClass::List.to_specialized_instance(self.db(), [Type::unknown()]) @@ -5303,95 +5288,167 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { elts, } = set; + let elts = elts.iter().map(|elt| [Some(elt)]); self.infer_collection_literal(elts, tcx, KnownClass::Set) .unwrap_or_else(|| { KnownClass::Set.to_specialized_instance(self.db(), [Type::unknown()]) }) } - // Infer the type of a collection literal expression. - fn infer_collection_literal( - &mut self, - elts: &[ast::Expr], - tcx: TypeContext<'db>, - collection_class: KnownClass, - ) -> Option> { - // Extract the type variable `T` from `list[T]` in typeshed. - fn elts_ty( - collection_class: KnownClass, - db: &dyn Db, - ) -> Option<(ClassLiteral<'_>, Type<'_>)> { - let class_literal = collection_class.try_to_class_literal(db)?; - let generic_context = class_literal.generic_context(db)?; - let variables = generic_context.variables(db); - let elts_ty = variables.iter().exactly_one().ok()?; - Some((class_literal, Type::TypeVar(*elts_ty))) - } - - let annotated_elts_ty = tcx - .known_specialization(collection_class, self.db()) - .and_then(|specialization| specialization.types(self.db()).iter().exactly_one().ok()) - .copied(); - - let (class_literal, elts_ty) = elts_ty(collection_class, self.db()).unwrap_or_else(|| { - let name = collection_class.name(self.db()); - panic!("Typeshed should always have a `{name}` class in `builtins.pyi` with a single type variable") - }); - - // Create a set of constraints to infer a precise type for `T`. - let mut builder = SpecializationBuilder::new(self.db()); - - match annotated_elts_ty { - // The annotated type acts as a constraint for `T`. - // - // Note that we infer the annotated type _before_ the elements, to closer match the order - // of any unions written in the type annotation. - Some(annotated_elts_ty) => { - builder.infer(elts_ty, annotated_elts_ty).ok()?; - } - - // If a valid type annotation was not provided, avoid restricting the type of the collection - // by unioning the inferred type with `Unknown`. - None => builder.infer(elts_ty, Type::unknown()).ok()?, - } - - // The inferred type of each element acts as an additional constraint on `T`. - for elt in elts { - let inferred_elt_ty = self.infer_expression(elt, TypeContext::new(annotated_elts_ty)); - - // Convert any element literals to their promoted type form to avoid excessively large - // unions for large nested list literals, which the constraint solver struggles with. - let inferred_elt_ty = - inferred_elt_ty.apply_type_mapping(self.db(), &TypeMapping::PromoteLiterals); - builder.infer(elts_ty, inferred_elt_ty).ok()?; - } - - let class_type = class_literal - .apply_specialization(self.db(), |generic_context| builder.build(generic_context)); - - Type::from(class_type).to_instance(self.db()) - } - - fn infer_dict_expression(&mut self, dict: &ast::ExprDict, _tcx: TypeContext<'db>) -> Type<'db> { + fn infer_dict_expression(&mut self, dict: &ast::ExprDict, tcx: TypeContext<'db>) -> Type<'db> { let ast::ExprDict { range: _, node_index: _, items, } = dict; - // TODO: Use the type context for more precise inference. - for item in items { - self.infer_optional_expression(item.key.as_ref(), TypeContext::default()); - self.infer_expression(&item.value, TypeContext::default()); + // Validate `TypedDict` dictionary literal assignments. + if let Some(typed_dict) = tcx.annotation.and_then(Type::into_typed_dict) { + let typed_dict_items = typed_dict.items(self.db()); + + for item in items { + self.infer_optional_expression(item.key.as_ref(), TypeContext::default()); + + if let Some(ast::Expr::StringLiteral(ref key)) = item.key + && let Some(key) = key.as_single_part_string() + && let Some(field) = typed_dict_items.get(key.as_str()) + { + self.infer_expression(&item.value, TypeContext::new(Some(field.declared_ty))); + } else { + self.infer_expression(&item.value, TypeContext::default()); + } + } + + validate_typed_dict_dict_literal( + &self.context, + typed_dict, + dict, + dict.into(), + |expr| self.expression_type(expr), + ); + + return Type::TypedDict(typed_dict); } - KnownClass::Dict.to_specialized_instance( - self.db(), - [ - todo_type!("dict literal key type"), - todo_type!("dict literal value type"), - ], - ) + // Avoid false positives for the functional `TypedDict` form, which is currently + // unsupported. + if let Some(Type::Dynamic(DynamicType::Todo(_))) = tcx.annotation { + return KnownClass::Dict + .to_specialized_instance(self.db(), [Type::unknown(), Type::unknown()]); + } + + let items = items + .iter() + .map(|item| [item.key.as_ref(), Some(&item.value)]); + + self.infer_collection_literal(items, tcx, KnownClass::Dict) + .unwrap_or_else(|| { + KnownClass::Dict + .to_specialized_instance(self.db(), [Type::unknown(), Type::unknown()]) + }) + } + + // Infer the type of a collection literal expression. + fn infer_collection_literal<'expr, const N: usize>( + &mut self, + elts: impl Iterator; N]>, + tcx: TypeContext<'db>, + collection_class: KnownClass, + ) -> Option> { + // Extract the type variable `T` from `list[T]` in typeshed. + fn elt_tys( + collection_class: KnownClass, + db: &dyn Db, + ) -> Option<(ClassLiteral<'_>, &FxOrderSet>)> { + let class_literal = collection_class.try_to_class_literal(db)?; + let generic_context = class_literal.generic_context(db)?; + Some((class_literal, generic_context.variables(db))) + } + + let (class_literal, elt_tys) = elt_tys(collection_class, self.db()).unwrap_or_else(|| { + let name = collection_class.name(self.db()); + panic!("Typeshed should always have a `{name}` class in `builtins.pyi`") + }); + + // Extract the annotated type of `T`, if provided. + let annotated_elt_tys = tcx + .known_specialization(collection_class, self.db()) + .map(|specialization| specialization.types(self.db())); + + // Create a set of constraints to infer a precise type for `T`. + let mut builder = SpecializationBuilder::new(self.db()); + + match annotated_elt_tys { + // The annotated type acts as a constraint for `T`. + // + // Note that we infer the annotated type _before_ the elements, to more closely match the + // order of any unions as written in the type annotation. + Some(annotated_elt_tys) => { + for (elt_ty, annotated_elt_ty) in iter::zip(elt_tys, annotated_elt_tys) { + builder + .infer(Type::TypeVar(*elt_ty), *annotated_elt_ty) + .ok()?; + } + } + + // If a valid type annotation was not provided, avoid restricting the type of the collection + // by unioning the inferred type with `Unknown`. + None => { + for elt_ty in elt_tys { + builder + .infer(Type::TypeVar(*elt_ty), Type::unknown()) + .ok()?; + } + } + } + + let elt_tcxs = match annotated_elt_tys { + None => Either::Left(iter::repeat(TypeContext::default())), + Some(tys) => Either::Right(tys.iter().map(|ty| TypeContext::new(Some(*ty)))), + }; + + for elts in elts { + // An unpacking expression for a dictionary. + if let &[None, Some(value)] = elts.as_slice() { + let inferred_value_ty = self.infer_expression(value, TypeContext::default()); + + // Merge the inferred type of the nested dictionary. + if let Some(specialization) = + inferred_value_ty.known_specialization(KnownClass::Dict, self.db()) + { + for (elt_ty, inferred_elt_ty) in + iter::zip(elt_tys, specialization.types(self.db())) + { + builder + .infer(Type::TypeVar(*elt_ty), *inferred_elt_ty) + .ok()?; + } + } + + continue; + } + + // The inferred type of each element acts as an additional constraint on `T`. + for (elt, elt_ty, elt_tcx) in itertools::izip!(elts, elt_tys, elt_tcxs.clone()) { + let Some(inferred_elt_ty) = self.infer_optional_expression(elt, elt_tcx) else { + continue; + }; + + // Convert any element literals to their promoted type form to avoid excessively large + // unions for large nested list literals, which the constraint solver struggles with. + let inferred_elt_ty = + inferred_elt_ty.apply_type_mapping(self.db(), &TypeMapping::PromoteLiterals); + + builder + .infer(Type::TypeVar(*elt_ty), inferred_elt_ty) + .ok()?; + } + } + + let class_type = class_literal + .apply_specialization(self.db(), |generic_context| builder.build(generic_context)); + + Type::from(class_type).to_instance(self.db()) } /// Infer the type of the `iter` expression of the first comprehension.