diff --git a/Cargo.lock b/Cargo.lock index a5956f32d427d..ab70cce0692e3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4865,6 +4865,7 @@ dependencies = [ "strum", "strum_macros", "test-case", + "thin-vec", "thiserror 2.0.19", "tracing", "ty_module_resolver", diff --git a/crates/ty_python_semantic/Cargo.toml b/crates/ty_python_semantic/Cargo.toml index ebdbfb5af864e..c8f1f5900379c 100644 --- a/crates/ty_python_semantic/Cargo.toml +++ b/crates/ty_python_semantic/Cargo.toml @@ -32,7 +32,7 @@ bitflags = { workspace = true } char_str = { workspace = true } compact_str = { workspace = true } drop_bomb = { workspace = true } -get-size2 = { workspace = true, features = ["indexmap", "ordermap"] } +get-size2 = { workspace = true, features = ["indexmap", "ordermap", "thin-vec"] } indexmap = { workspace = true } itertools = { workspace = true } memchr = { workspace = true } @@ -47,6 +47,7 @@ static_assertions = { workspace = true } strum = { workspace = true } strum_macros = { workspace = true } thiserror = { workspace = true } +thin-vec = { workspace = true } tracing = { workspace = true } [dev-dependencies] diff --git a/crates/ty_python_semantic/resources/mdtest/cycle.md b/crates/ty_python_semantic/resources/mdtest/cycle.md index 0e808c34775f5..49993a608b004 100644 --- a/crates/ty_python_semantic/resources/mdtest/cycle.md +++ b/crates/ty_python_semantic/resources/mdtest/cycle.md @@ -385,6 +385,80 @@ class Cached: reveal_type(Cached().metadata) # revealed: int ``` +## Guarded instance attributes when the base is checked first + +A guarded bound-method initializer remains valid, while another initializer still reports an +attribute that is missing from the base class. + +`base.py`: + +```py +class Base: + def __init__(self): + if not hasattr(self, "x"): + self.x = self.__str__ + if not hasattr(self, "z"): + self.z = self.y # error: [unresolved-attribute] +``` + +`child.py`: + +```py +from base import Base + +class Child(Base): + x = Base.__str__ + + def z(self): ... + def y(self): ... +``` + +## Guarded instance attributes when the subclass is checked first + +Checking the subclass first preserves the valid initializer and the missing-attribute diagnostic. + +`child.py`: + +```py +from base import Base + +class Child(Base): + x = Base.__str__ + + def z(self): ... + def y(self): ... +``` + +`base.py`: + +```py +class Base: + def __init__(self): + if not hasattr(self, "x"): + self.x = self.__str__ + if not hasattr(self, "z"): + self.z = self.y # error: [unresolved-attribute] +``` + +## Assignments in the opposite guard branch do not initialize an attribute + +Assigning an existing attribute when `hasattr` succeeds does not initialize it in the opposite +branch. That branch remains unreachable and cannot create another instance attribute. + +```py +class C: + def __init__(self): + self.x = 1 + + def update(self): + if hasattr(self, "x"): + self.x = 2 + else: + self.y = self.missing + +C().y # error: [unresolved-attribute] +``` + ## Decorator defined on a base class with constrained typevars, accessed from a subclass with decorated generic parameters This example was minimized from diff --git a/crates/ty_python_semantic/resources/mdtest/typed_dict.md b/crates/ty_python_semantic/resources/mdtest/typed_dict.md index 8f6feb0077ea8..59ed5a5304f77 100644 --- a/crates/ty_python_semantic/resources/mdtest/typed_dict.md +++ b/crates/ty_python_semantic/resources/mdtest/typed_dict.md @@ -2653,8 +2653,9 @@ def _(item: Item | str) -> None: ``` A successful membership test for an undeclared key narrows each union member to an intersection with -a synthesized `TypedDict`. Its mapping methods should retain their precise types, and copying the -narrowed union should remain efficient even when each member has a distinct optional field: +a synthesized protocol that records membership and subscript access. Its mapping methods should +retain their precise types, and copying the narrowed union should remain efficient even when each +member has a distinct optional field: ```py from typing import NotRequired @@ -2693,7 +2694,7 @@ def _(item: MembershipItem) -> None: def _(item: MembershipA) -> None: if "missing" in item: - reveal_type(item.copy()) # revealed: MembershipA & + reveal_type(item.copy()) # revealed: MembershipA & ``` Adding a regular dictionary to the union should not make copying it slow: @@ -5180,7 +5181,7 @@ def _(p: Person) -> None: reveal_type(p.setdefault("name", "Alice")) # revealed: str # __contains__ - reveal_type("name" in p) # revealed: bool + reveal_type("name" in p) # revealed: Literal[True] # __setitem__ p["name"] = "Alice" @@ -6109,7 +6110,8 @@ the key, because "extra items" are allowed by default. For example, even though define a `"foo"` field, it could be _assigned to_ with another `TypedDict` that does: ```py -from typing_extensions import Literal +from collections.abc import Mapping +from typing_extensions import Final, Literal, NotRequired, TypeGuard class Foo(TypedDict): foo: int @@ -6120,14 +6122,14 @@ class Bar(TypedDict): def disappointment(u: Foo | Bar, v: Literal["foo"]): if "foo" in u: # We don't narrow to just `Foo` here... - reveal_type(u) # revealed: Foo | (Bar & ) + reveal_type(u) # revealed: Foo | (Bar & ) reveal_type(u["foo"]) # revealed: object else: # ...(even though we *can* narrow it here)... reveal_type(u) # revealed: Bar if v in u: - reveal_type(u) # revealed: Foo | (Bar & ) + reveal_type(u) # revealed: Foo | (Bar & ) reveal_type(u["foo"]) # revealed: object else: reveal_type(u) # revealed: Bar @@ -6139,16 +6141,124 @@ class FooBar(TypedDict): static_assert(is_assignable_to(FooBar, Foo)) static_assert(is_assignable_to(FooBar, Bar)) +``` + +A successful membership check permits subscript access even when the key is not included in the +annotated key type: +```py def dictionary_union(u: Foo | dict[Literal["a", "b"], int]): if "c" in u: - # TODO: This should stop erroring if we prove that the `dict` arm cannot contain `"c"`. - # error: [invalid-argument-type] reveal_type(u["c"]) # revealed: object +def mapping_union(u: Foo | Mapping[Literal["a", "b"], int]): + if "c" in u: + reveal_type(u["c"]) # revealed: object + +def mapping_membership(mapping: Mapping[Literal["a", "b"], int]): + if "c" in mapping: + reveal_type(mapping["c"]) # revealed: object +``` + +When a condition checks multiple keys, each successful check is retained: + +```py +def combined_typed_dict_checks(u: Foo | Bar): + has_foo = "foo" in u + has_bar = "bar" in u + if has_foo and has_bar: + reveal_type(u["foo"]) # revealed: object + reveal_type(u["bar"]) # revealed: object + +def combined_mapping_checks(mapping: Mapping[Literal["a", "b"], int]): + has_c = "c" in mapping + has_d = "d" in mapping + if has_c and has_d: + reveal_type(mapping["c"]) # revealed: object + reveal_type(mapping["d"]) # revealed: object + +def either_key_is_present(mapping: Mapping[Literal["a"], int]): + if "c" in mapping or "d" in mapping: + if "c" not in mapping: + reveal_type(mapping["d"]) # revealed: object +``` + +Membership checks that occur after a `TypeGuard` still apply to the replacement type. However, a +`TypeGuard` discards membership facts from preceding conditions, along with all other previously +known type information: + +```py +def guard_object(value: object) -> TypeGuard[object]: + return True + +def guard_bar(value: object) -> TypeGuard[Bar]: + return True + +def membership_after_typeguard(value: Foo | Literal["abc"]): + has_z = "z" in value + if guard_bar(value) and has_z: + reveal_type(value["z"]) # revealed: object + if has_z and guard_bar(value): + value["z"] # error: [invalid-key] + +class OptionalKey(TypedDict): + x: NotRequired[int] + +def optional_key_after_typeguard(value: OptionalKey): + has_x = "x" in value + if guard_bar(value) and has_x: + reveal_type(value["x"]) # revealed: object + +class AlwaysContains(Mapping[str, int]): + def __contains__(self, key: object, /) -> Literal[True]: + return True + +class SometimesContains(Mapping[str, int]): ... +class Target: ... + +def guard_mapping_or_target(value: object) -> TypeGuard[AlwaysContains | Target]: + return True + +def absent_key_after_typeguard(value: SometimesContains): + lacks_x = "x" not in value + if guard_mapping_or_target(value) and lacks_x: + reveal_type(value) # revealed: Target + if lacks_x and guard_mapping_or_target(value): + reveal_type(value) # revealed: AlwaysContains | Target + +def mapping_membership_after_typeguard(u: Foo | Mapping[Literal["a", "b"], int]): + has_c = "c" in u + if guard_object(u) and has_c: + reveal_type(u["c"]) # revealed: object +``` + +Precomputed refinements, such as filtering a union by a nominal tag, preserve preceding membership +facts: + +```py +class UserSettings(Mapping[Literal["user_id"], int]): + kind: Final[Literal["user"]] = "user" + +class SystemSettings(Mapping[Literal["system_id"], int]): + kind: Final[Literal["system"]] = "system" + +def read_timeout(settings: UserSettings | SystemSettings) -> object: + has_timeout = "timeout" in settings + is_user_settings = settings.kind == "user" + + if has_timeout and is_user_settings: + return settings["timeout"] + + raise KeyError("timeout") +``` + +For other objects, a successful membership check does not imply that the same value can be used as a +subscript: + +```py def literal_union(u: Foo | Literal["abc"]): if "a" in u: - # revealed: (Foo & ) | (Literal["abc"] & ) + # revealed: (Foo & ) | (Literal["abc"] & ) reveal_type(u) def literal_union_key_access(obj: Foo | Literal["a"]): @@ -6198,9 +6308,8 @@ def _(t: Bar, u: Foo | Intersection[Bar, Any], v: Intersection[Bar, Any], w: Lit if "bar" not in u: reveal_type(u) # revealed: Foo else: - # TODO: This should simplify to `Foo | (Bar & Any)`, since `Foo` is a - # subtype of the synthesized protocol. - reveal_type(u) # revealed: (Foo & ) | (Bar & Any) + # `Foo` is open, so it may contain an undeclared `"bar"` key. + reveal_type(u) # revealed: (Foo & ) | (Bar & Any) if "bar" not in v: reveal_type(v) # revealed: Never @@ -6210,12 +6319,12 @@ def _(t: Bar, u: Foo | Intersection[Bar, Any], v: Intersection[Bar, Any], w: Lit if w not in u: reveal_type(u) # revealed: Foo else: - reveal_type(u) # revealed: (Foo & ) | (Bar & Any) + reveal_type(u) # revealed: (Foo & ) | (Bar & Any) if "bar" not in (u2 := u): reveal_type(u2) # revealed: Foo else: - reveal_type(u2) # revealed: (Foo & ) | (Bar & Any) + reveal_type(u2) # revealed: (Foo & ) | (Bar & Any) ``` With `closed=True`, the narrowing that we couldn't do above becomes possible, because a [closed] @@ -6479,7 +6588,7 @@ def test_in(x: ThingWithBaz): if "baz" not in x: reveal_type(x) # revealed: Foo else: - reveal_type(x) # revealed: (Foo & ) | Baz + reveal_type(x) # revealed: (Foo & ) | Baz ``` Nested PEP 695 type aliases (an alias referring to another alias) also work: @@ -6508,7 +6617,7 @@ def test_nested_in(x: OuterWithBaz): if "baz" not in x: reveal_type(x) # revealed: Foo else: - reveal_type(x) # revealed: (Foo & ) | Baz + reveal_type(x) # revealed: (Foo & ) | Baz ``` ## Only annotated declarations are allowed in the class body diff --git a/crates/ty_python_semantic/src/types/class/typed_dict.rs b/crates/ty_python_semantic/src/types/class/typed_dict.rs index fa2a1586fc61e..f95c690cb5069 100644 --- a/crates/ty_python_semantic/src/types/class/typed_dict.rs +++ b/crates/ty_python_semantic/src/types/class/typed_dict.rs @@ -39,6 +39,12 @@ pub(super) fn synthesize_typed_dict_method<'db>( let instance_ty = Type::TypedDict(typed_dict); match method_name { "__init__" => Some(synthesize_typed_dict_init(db, env, typed_dict, fields())), + "__contains__" => Some(synthesize_typed_dict_contains( + db, + env, + instance_ty, + fields(), + )), "__getitem__" => Some(synthesize_typed_dict_getitem(db, env, typed_dict, fields())), "__setitem__" => Some(synthesize_typed_dict_setitem(db, env, typed_dict, fields())), "__delitem__" => Some(synthesize_typed_dict_delitem(db, env, typed_dict, fields())), @@ -127,6 +133,43 @@ impl<'db> TypedDictFields<'db> { } } +/// Synthesize the `__contains__` method for a `TypedDict`. +fn synthesize_typed_dict_contains<'db>( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + instance_ty: Type<'db>, + fields: TypedDictFields<'db>, +) -> Type<'db> { + let overloads = fields + .iter() + .filter(|(_, field)| field.is_required()) + .map(|(field_name, _)| { + let parameters = [ + Parameter::positional_only(Some(Name::new_static("self"))) + .with_annotated_type(instance_ty), + Parameter::positional_only(Some(Name::new_static("key"))) + .with_annotated_type(Type::string_literal(db, field_name)), + ]; + Signature::new(Parameters::standard(parameters), Type::bool_literal(true)) + }) + .chain(std::iter::once(Signature::new( + Parameters::standard([ + Parameter::positional_only(Some(Name::new_static("self"))) + .with_annotated_type(instance_ty), + Parameter::positional_only(Some(Name::new_static("key"))) + .with_annotated_type(Type::object()), + ]), + KnownClass::Bool.to_instance(db, env), + ))); + + Type::Callable(CallableType::new( + db, + CallableSignature::from_overloads(overloads), + CallableTypeKind::FunctionLike, + CallableFunctionProvenance::None, + )) +} + /// Synthesize the `__init__` method for a `TypedDict`. /// /// overloads: @@ -1027,7 +1070,7 @@ impl<'db> DynamicTypedDictLiteral<'db> { return member.inner; } - // Fall back to TypedDictFallback for methods like __contains__, items, keys, etc. + // Fall back to TypedDictFallback for methods like items, keys, etc. // This mirrors the behavior of StaticClassLiteral::typed_dict_member. typed_dict_class_member( db, diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index cc0956b59c511..a9541295f7be3 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -3,22 +3,21 @@ use std::collections::{BTreeMap, btree_map::Entry as BTreeEntry, hash_map::Entry use crate::reachability::{narrow_type_by_constraint, type_narrowed_by_previous_patterns}; use crate::subscript::PyIndex; +use crate::types::call::CallArguments; use crate::types::function::KnownFunction; use crate::types::infer::{ExpressionInference, infer_same_file_expression_type}; use crate::types::special_form::TypeQualifier; use crate::types::tuple::{TupleLength, TupleSpec, TupleSpecBuilder, TupleType, TupleUnpacker}; -use crate::types::typed_dict::{ - TypedDictField, TypedDictFieldBuilder, TypedDictSchema, TypedDictType, -}; +use crate::types::typed_dict::{TypedDictFieldBuilder, TypedDictSchema, TypedDictType}; use crate::types::{ CallableType, ClassBase, ClassLiteral, ClassPatternPositionalSource, ClassType, IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType, LiteralValueTypeKind, - Parameter, Parameters, Signature, SpecialFormType, SubclassOfInner, SubclassOfType, Truthiness, - Type, TypeContext, TypeVarBoundOrConstraints, UnionBuilder, callable_pattern_type, - class_pattern_positional_sources, definite_match_pattern_type_for_subject, - exact_sequence_pattern_type, infer_expression_types, mapping_pattern_type, - pattern_binding_fallthrough_type, sequence_pattern_type_builder, singleton_pattern_type, - starred_sequence_pattern_type, typed_dict_matches_class_pattern, + Parameter, Parameters, Signature, SpecialFormType, StringLiteralType, SubclassOfInner, + SubclassOfType, Truthiness, Type, TypeContext, TypeVarBoundOrConstraints, UnionBuilder, + callable_pattern_type, class_pattern_positional_sources, + definite_match_pattern_type_for_subject, exact_sequence_pattern_type, infer_expression_types, + mapping_pattern_type, pattern_binding_fallthrough_type, sequence_pattern_type_builder, + singleton_pattern_type, starred_sequence_pattern_type, typed_dict_matches_class_pattern, }; use crate::{Db, ProgramEnvironment}; use ty_python_core::expression::Expression; @@ -29,6 +28,7 @@ use ty_python_core::predicate::{ PatternPredicateKind, Predicate, PredicateNode, SequencePatternPredicateKind, SubjectElementPatternPredicate, }; +use ty_python_core::reachability_constraints::ScopedReachabilityConstraintId; use ty_python_core::scope::ScopeId; use ty_python_core::{ExpressionNodeKey, NarrowingEvaluator, place_table, semantic_index}; @@ -37,7 +37,6 @@ use ruff_python_ast::name::Name; use ruff_python_stdlib::identifiers::is_identifier; use super::UnionType; -use super::call::CallArguments; use super::constraints::{ConstraintSetBuilder, PathBounds, Solutions}; use super::equality::{ ComparisonSoundnessPolicy, equality_exclusion_constraint, equality_truthiness, @@ -50,6 +49,7 @@ use ruff_python_ast as ast; use ruff_python_ast::{BoolOp, ExprBoolOp}; use rustc_hash::FxHashMap; use smallvec::{SmallVec, smallvec, smallvec_inline}; +use thin_vec::{ThinVec, thin_vec}; mod containment; @@ -712,61 +712,250 @@ impl<'db> NarrowingOperation<'db> { } } +/// One conjunction in the DNF representation. +/// +/// Key membership checks are stored as facts instead of converted into types immediately. For +/// `value: A | B`, converting each check in `"a" in value and "b" in value` can split both checks +/// into `A` and `B`; combining them would then create every pairing of those branches. Keeping the +/// checks here applies them only when the final type is evaluated. It also lets a key fact constrain +/// a later precomputed refinement instead of restoring the original type. +#[derive(Hash, PartialEq, Debug, Eq, Clone, get_size2::GetSize, salsa::SalsaValue)] +struct ConstraintConjunction<'db> { + type_constraints: SmallVec<[NarrowingOperation<'db>; 2]>, + key_constraints: ThinVec>, +} + #[derive(Hash, PartialEq, Debug, Eq, Clone, get_size2::GetSize, salsa::SalsaValue)] -struct Conjunctions<'db> { - conjuncts: SmallVec<[NarrowingOperation<'db>; 2]>, +enum KeyConstraint<'db> { + Present(PresentKeyConstraint<'db>), + Absent(AbsentKeyConstraint<'db>), } -impl<'db> Conjunctions<'db> { +/// A deferred fact that `key` is present in `source`. +#[derive(Hash, PartialEq, Debug, Eq, Clone, get_size2::GetSize, salsa::SalsaValue)] +struct PresentKeyConstraint<'db> { + source: Type<'db>, + key: StringLiteralType<'db>, +} + +impl<'db> PresentKeyConstraint<'db> { + /// Apply this fact without restoring type information discarded by a replacement narrowing. + /// + /// `source` determines what the original membership test proved, but any resulting key + /// constraint is applied to `replacement`: + /// + /// ```python + /// has_x = "x" in value + /// if guard_to_bar(value) and has_x: + /// value["x"] # The key fact applies to `Bar`, not the original type. + /// ``` + fn apply_after_replacement( + &self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + replacement: Type<'db>, + ) -> Type<'db> { + if key_is_always_present(db, env, self.source, self.key) { + return replacement; + } + if narrow_with_present_key(db, env, self.source, self.key.value(db)).is_never() { + return Type::Never; + } + + if key_membership_implies_subscript(db, env, self.source) { + IntersectionType::from_two_elements( + db, + env, + replacement, + mapping_present_key_protocol(db, env, self.key.value(db)), + ) + } else { + narrow_with_present_key(db, env, replacement, self.key.value(db)) + } + } + + fn apply( + &self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + ty: Type<'db>, + is_replacement: bool, + ) -> Type<'db> { + if is_replacement { + self.apply_after_replacement(db, env, ty) + } else { + narrow_with_present_key(db, env, ty, self.key.value(db)) + } + } +} + +/// A deferred key-absence operation. Keeping the filtered source out of the conjunction prevents +/// it from replacing a preceding `TypeGuard` result. +/// +/// ```python +/// lacks_x = "x" not in value +/// if guard_to_bar(value) and lacks_x: +/// reveal_type(value) # Bar, filtered by the key-absence fact +/// ``` +#[derive(Hash, PartialEq, Debug, Eq, Clone, get_size2::GetSize, salsa::SalsaValue)] +struct AbsentKeyConstraint<'db> { + source: Type<'db>, + key: StringLiteralType<'db>, +} + +impl<'db> AbsentKeyConstraint<'db> { + /// Apply this fact to `replacement`, using `source` only to detect an impossible condition. + fn apply_after_replacement( + &self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + replacement: Type<'db>, + ) -> Type<'db> { + if narrow_with_absent_key(db, env, self.source, self.key).is_never() { + Type::Never + } else { + narrow_with_absent_key(db, env, replacement, self.key) + } + } + + fn apply( + &self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + ty: Type<'db>, + is_replacement: bool, + ) -> Type<'db> { + if is_replacement { + self.apply_after_replacement(db, env, ty) + } else { + narrow_with_absent_key(db, env, ty, self.key) + } + } +} + +impl<'db> ConstraintConjunction<'db> { fn singleton(ty: Type<'db>) -> Self { Self { - conjuncts: smallvec![NarrowingOperation::Intersection(ty)], + type_constraints: smallvec![NarrowingOperation::Intersection(ty)], + key_constraints: thin_vec![], } } fn generic_filtering(ty: Type<'db>) -> Self { Self { - conjuncts: smallvec![NarrowingOperation::GenericFiltering(ty)], + type_constraints: smallvec![NarrowingOperation::GenericFiltering(ty)], + key_constraints: thin_vec![], + } + } + + fn present_key(source: Type<'db>, key: StringLiteralType<'db>) -> Self { + Self { + type_constraints: smallvec![], + key_constraints: thin_vec![KeyConstraint::Present(PresentKeyConstraint { + source, + key, + })], + } + } + + fn absent_key(source: Type<'db>, key: StringLiteralType<'db>) -> Self { + Self { + type_constraints: smallvec![], + key_constraints: thin_vec![KeyConstraint::Absent(AbsentKeyConstraint { source, key })], } } fn and_with(mut self, other: Self) -> Self { if self - .conjuncts + .type_constraints .iter() .any(|conjunct| conjunct.ty().is_never()) || other - .conjuncts + .type_constraints .iter() .any(|conjunct| conjunct.ty().is_never()) { return Self::singleton(Type::Never); } - for conjunct in other.conjuncts { - if !self.conjuncts.contains(&conjunct) { - self.conjuncts.push(conjunct); + for type_constraint in other.type_constraints { + if !self.type_constraints.contains(&type_constraint) { + self.type_constraints.push(type_constraint); + } + } + + if self.key_constraints.is_empty() { + self.key_constraints = other.key_constraints; + } else { + for key_constraint in other.key_constraints { + if !self.key_constraints.contains(&key_constraint) { + self.key_constraints.push(key_constraint); + } + } + } + + self + } + + /// Add only the deferred key facts from `other`, leaving its type constraints behind. + fn and_with_key_constraints_from(mut self, other: &Self) -> Self { + if self.key_constraints.is_empty() { + self.key_constraints = other.key_constraints.clone(); + } else { + for key_constraint in &other.key_constraints { + if !self.key_constraints.contains(key_constraint) { + self.key_constraints.push(key_constraint.clone()); + } } } + self } - fn evaluate_constraint_type(self, db: &'db dyn Db, env: &ProgramEnvironment<'db>) -> Type<'db> { - if self.conjuncts.len() == 1 { - return self.conjuncts[0].ty(); + /// Materialize this conjunction, applying deferred key facts after ordinary type constraints. + /// + /// Present-key facts are applied before absent-key facts so contradictory facts for the same + /// key evaluate to `Never`. Replacement disjuncts use the source-aware application methods to + /// avoid undoing `TypeGuard` narrowing. + fn evaluate_constraint_type( + self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + is_replacement: bool, + ) -> Type<'db> { + if self.key_constraints.is_empty() && self.type_constraints.len() == 1 { + return self.type_constraints[0].ty(); } // Collapse shared union arms before distributing the next constraint over them. - self.conjuncts - .into_iter() - .fold(Type::object(), |accumulated, conjunct| match conjunct { - NarrowingOperation::Intersection(ty) => { - IntersectionType::from_two_elements(db, env, accumulated, ty) - } - NarrowingOperation::GenericFiltering(ty) => { - filter_generic_narrowing_constraint(db, env, accumulated, ty) - } - }) + let mut current = + self.type_constraints + .into_iter() + .fold(Type::object(), |accumulated, conjunct| match conjunct { + NarrowingOperation::Intersection(ty) => { + IntersectionType::from_two_elements(db, env, accumulated, ty) + } + NarrowingOperation::GenericFiltering(ty) => { + filter_generic_narrowing_constraint(db, env, accumulated, ty) + } + }); + + if self.key_constraints.is_empty() { + return current; + } + + for key_constraint in &self.key_constraints { + if let KeyConstraint::Present(present_key) = key_constraint { + current = present_key.apply(db, env, current, is_replacement); + } + } + for key_constraint in &self.key_constraints { + if let KeyConstraint::Absent(absent_key) = key_constraint { + current = absent_key.apply(db, env, current, is_replacement); + } + } + + current } } @@ -1041,64 +1230,101 @@ fn specialize_generic_class_from_solutions<'db>( /// Represents narrowing constraints in Disjunctive Normal Form (DNF). /// /// This is a disjunction (OR) of conjunctions (AND) of constraints. -/// The DNF representation allows us to properly track "replacement" constraints -/// (created by `TypeGuard` types and similar) through boolean operations. +/// The DNF representation allows us to distinguish intersections, precomputed refinements, and +/// `TypeGuard` replacements through boolean operations. /// /// For example: /// - `f(x) and g(x)` where f returns `TypeIs[A]` and g returns `TypeGuard[B]` /// => and -/// ===> `NarrowingConstraint { intersection_disjuncts: [A], replacement_disjuncts: [] }` -/// ===> `NarrowingConstraint { intersection_disjuncts: [], replacement_disjuncts: [B] }` -/// => `NarrowingConstraint { intersection_disjuncts: [], replacement_disjuncts: [B] }` +/// ===> an intersection disjunct for `A` +/// ===> a replacement disjunct for `B` +/// => a replacement disjunct for `B` /// => evaluates to `B` (`TypeGuard` clobbers any previous type information) /// /// - `f(x) or g(x)` where f returns `TypeIs[A]` and g returns `TypeGuard[B]` /// => or -/// ===> `NarrowingConstraint { intersection_disjuncts: [A], replacement_disjuncts: [] }` -/// ===> `NarrowingConstraint { intersection_disjuncts: [], replacement_disjuncts: [B] }` -/// => `NarrowingConstraint { intersection_disjuncts: [A], replacement_disjuncts: [B] }` +/// ===> an intersection disjunct for `A` +/// ===> a replacement disjunct for `B` +/// => both disjuncts are retained /// => evaluates to `(P & A) | B`, where `P` is our previously-known type #[derive(Hash, PartialEq, Debug, Eq, Clone, get_size2::GetSize, salsa::SalsaValue)] pub(crate) struct NarrowingConstraint<'db> { - /// Intersection constraint (from `isinstance()` narrowing comparisons, `TypeIs`, and - /// similar). We keep these as a disjunction of conjunctions to avoid constructing - /// union/intersection types while merging constraints. - intersection_disjuncts: SmallVec<[Conjunctions<'db>; 1]>, - - /// "Replacement" constraints: instead of intersecting the previous type with a new type, - /// the previous type is simply replaced wholesale with the new type. A common use case for - /// these constraints is `typing.TypeGuard`. We can't eagerly union disjunctions because - /// `TypeGuard` clobbers the previously-known type; within each replacement disjunct, however, - /// we may eagerly intersect conjunctions with a later intersection narrowing. - replacement_disjuncts: SmallVec<[Conjunctions<'db>; 1]>, + disjuncts: SmallVec<[ConstraintDisjunct<'db>; 1]>, +} + +#[derive(Hash, PartialEq, Debug, Eq, Clone, get_size2::GetSize, salsa::SalsaValue)] +struct ConstraintDisjunct<'db> { + kind: ConstraintDisjunctKind, + conjunction: ConstraintConjunction<'db>, +} + +#[derive(Hash, PartialEq, Debug, Eq, Clone, Copy, get_size2::GetSize, salsa::SalsaValue)] +enum ConstraintDisjunctKind { + /// Intersect the previous type with this disjunct. + Intersection, + + /// Replace the previous type with a directly filtered refinement while retaining deferred key + /// facts. This is used when narrowing cannot be represented as an intersection. + PrecomputedRefinement, + + /// Replace all previously known type information with this disjunct. + TypeGuardReplacement, } impl<'db> NarrowingConstraint<'db> { + fn singleton(kind: ConstraintDisjunctKind, conjunction: ConstraintConjunction<'db>) -> Self { + Self { + disjuncts: smallvec_inline![ConstraintDisjunct { kind, conjunction }], + } + } + /// Create an "intersection" constraint: the previous type will be /// intersected with this constraint pub(crate) fn intersection(constraint: Type<'db>) -> Self { - Self { - intersection_disjuncts: smallvec_inline![Conjunctions::singleton(constraint)], - replacement_disjuncts: smallvec![], - } + Self::singleton( + ConstraintDisjunctKind::Intersection, + ConstraintConjunction::singleton(constraint), + ) } /// Create an intersection constraint that preserves generic arguments already known about /// the subject when narrowing it to a subclass. fn generic_filtering(constraint: Type<'db>) -> Self { - Self { - intersection_disjuncts: smallvec_inline![Conjunctions::generic_filtering(constraint)], - replacement_disjuncts: smallvec![], - } + Self::singleton( + ConstraintDisjunctKind::Intersection, + ConstraintConjunction::generic_filtering(constraint), + ) } - /// Create a "replacement" constraint: the previous type will be - /// replaced wholesale with this constraint - fn replacement(constraint: Type<'db>) -> Self { - Self { - intersection_disjuncts: smallvec![], - replacement_disjuncts: smallvec_inline![Conjunctions::singleton(constraint)], - } + /// Create a precomputed refinement: the previous type will be replaced with this constraint, + /// but preceding deferred key facts will still apply. + fn precomputed_refinement(constraint: Type<'db>) -> Self { + Self::singleton( + ConstraintDisjunctKind::PrecomputedRefinement, + ConstraintConjunction::singleton(constraint), + ) + } + + /// Create a `TypeGuard` replacement that discards all previously known type information. + fn type_guard_replacement(constraint: Type<'db>) -> Self { + Self::singleton( + ConstraintDisjunctKind::TypeGuardReplacement, + ConstraintConjunction::singleton(constraint), + ) + } + + fn present_key(source: Type<'db>, key: StringLiteralType<'db>) -> Self { + Self::singleton( + ConstraintDisjunctKind::Intersection, + ConstraintConjunction::present_key(source, key), + ) + } + + fn absent_key(source: Type<'db>, key: StringLiteralType<'db>) -> Self { + Self::singleton( + ConstraintDisjunctKind::Intersection, + ConstraintConjunction::absent_key(source, key), + ) } /// Merge two constraints, taking their intersection but respecting "replacement" semantics (with @@ -1107,59 +1333,79 @@ impl<'db> NarrowingConstraint<'db> { // Distribute AND over OR: (A1 | A2 | ...) AND (B1 | B2 | ...) // becomes (A1 & B1) | (A1 & B2) | ... | (A2 & B1) | ... // - // In our representation, the RHS `replacement_disjuncts` will all clobber the LHS disjuncts - // when they are `and`ed, so they'll just stay as is. + // RHS `TypeGuard` replacements discard the LHS disjuncts entirely. RHS precomputed + // refinements clobber the LHS type constraints, but facts established by LHS key checks + // still apply to the refined type. // - // The thing we actually need to deal with is the RHS `intersection_disjuncts`. Each RHS - // disjunct gets intersected with each LHS disjunct, producing the cartesian product. - // This is still deferred as conjunction lists. + // Each RHS intersection disjunct gets intersected with each LHS disjunct, producing the + // cartesian product. The merged disjunct retains the LHS kind. // - // We also intersect each LHS `replacement_disjunct` with every RHS intersection disjunct - // to form new additional `replacement_disjuncts`. - if other.intersection_disjuncts.is_empty() { - return other; - } - - let mut new_intersection_disjuncts = smallvec![]; - for intersection_disjunct in &self.intersection_disjuncts { - for other_intersection_disjunct in &other.intersection_disjuncts { - let merged = intersection_disjunct - .clone() - .and_with(other_intersection_disjunct.clone()); - if !new_intersection_disjuncts.contains(&merged) { - new_intersection_disjuncts.push(merged); + // The conjunctions still defer constructing the corresponding intersection types. + let mut other_intersection_disjuncts: SmallVec<[ConstraintConjunction<'db>; 1]> = + smallvec![]; + let mut other_precomputed_refinement_conjunctions: SmallVec< + [ConstraintConjunction<'db>; 1], + > = smallvec![]; + let mut new_disjuncts: SmallVec<[ConstraintDisjunct<'db>; 1]> = smallvec![]; + for disjunct in other.disjuncts { + match disjunct.kind { + ConstraintDisjunctKind::Intersection => { + other_intersection_disjuncts.push(disjunct.conjunction); + } + ConstraintDisjunctKind::PrecomputedRefinement => { + other_precomputed_refinement_conjunctions.push(disjunct.conjunction); + } + ConstraintDisjunctKind::TypeGuardReplacement => { + if !new_disjuncts.contains(&disjunct) { + new_disjuncts.push(disjunct); + } } } } - let mut additional_replacement_disjuncts: SmallVec<[Conjunctions<'db>; 1]> = smallvec![]; - for replacement_disjunct in &self.replacement_disjuncts { - for other_intersection_disjunct in &other.intersection_disjuncts { - let merged = replacement_disjunct - .clone() - .and_with(other_intersection_disjunct.clone()); - if !additional_replacement_disjuncts.contains(&merged) { - additional_replacement_disjuncts.push(merged); + for disjunct in &self.disjuncts { + for other_conjunction in &other_precomputed_refinement_conjunctions { + let merged = ConstraintDisjunct { + kind: ConstraintDisjunctKind::PrecomputedRefinement, + conjunction: other_conjunction + .clone() + .and_with_key_constraints_from(&disjunct.conjunction), + }; + if !new_disjuncts.contains(&merged) { + new_disjuncts.push(merged); } } } - let mut new_replacement_disjuncts = other.replacement_disjuncts; + if other_intersection_disjuncts.is_empty() { + return Self { + disjuncts: new_disjuncts, + }; + } - new_replacement_disjuncts.extend(additional_replacement_disjuncts); + for disjunct in &self.disjuncts { + for other_conjunction in &other_intersection_disjuncts { + let merged = ConstraintDisjunct { + kind: disjunct.kind, + conjunction: disjunct + .conjunction + .clone() + .and_with(other_conjunction.clone()), + }; + if !new_disjuncts.contains(&merged) { + new_disjuncts.push(merged); + } + } + } - NarrowingConstraint { - intersection_disjuncts: new_intersection_disjuncts, - replacement_disjuncts: new_replacement_disjuncts, + Self { + disjuncts: new_disjuncts, } } /// Merge two constraints with OR semantics (union/disjunction). fn merge_constraint_or(&mut self, other: Self) { - self.intersection_disjuncts - .extend(other.intersection_disjuncts); - self.replacement_disjuncts - .extend(other.replacement_disjuncts); + self.disjuncts.extend(other.disjuncts); } /// Evaluate the type this effectively constrains to @@ -1171,12 +1417,12 @@ impl<'db> NarrowingConstraint<'db> { env: &ProgramEnvironment<'db>, ) -> Type<'db> { let mut union = UnionBuilder::new(db, env); - for conjunctions in self - .replacement_disjuncts - .into_iter() - .chain(self.intersection_disjuncts) - { - union.add_in_place(conjunctions.evaluate_constraint_type(db, env)); + for disjunct in self.disjuncts { + union.add_in_place(disjunct.conjunction.evaluate_constraint_type( + db, + env, + disjunct.kind != ConstraintDisjunctKind::Intersection, + )); } union.build() } @@ -3193,6 +3439,59 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { place_table(db, self.scope()) } + /// Returns whether this predicate controls an assignment to the given member. + fn predicate_initializes_member(&self, member: ScopedPlaceId) -> bool { + let db = self.db; + let scope = self.scope(); + let index = semantic_index(db, scope.program_file(db)); + let use_def = index.use_def_map(scope.file_scope_id(db)); + let constraints = use_def.reachability_constraints(); + let predicates = use_def.predicates(); + let is_terminal = |constraint| { + matches!( + constraint, + ScopedReachabilityConstraintId::ALWAYS_TRUE + | ScopedReachabilityConstraintId::AMBIGUOUS + | ScopedReachabilityConstraintId::ALWAYS_FALSE + ) + }; + + let mut initialized_by_predicate = false; + for binding in use_def.reachable_bindings(member) { + if binding.binding.definition().is_none() { + continue; + } + + if binding.reachability_constraint == ScopedReachabilityConstraintId::ALWAYS_TRUE { + return false; + } + if is_terminal(binding.reachability_constraint) { + continue; + } + + let node = constraints.get_interior_node(binding.reachability_constraint); + let predicate = predicates[node.atom()]; + let opposing_branch = if self.is_positive == predicate.is_positive { + node.if_false() + } else { + node.if_true() + }; + + if predicate.node != self.predicate + || opposing_branch != ScopedReachabilityConstraintId::ALWAYS_FALSE + || ![node.if_true(), node.if_ambiguous(), node.if_false()] + .into_iter() + .all(is_terminal) + { + return false; + } + + initialized_by_predicate = true; + } + + initialized_by_predicate + } + fn scope(&self) -> ScopeId<'db> { let db = self.db; match self.predicate { @@ -3856,7 +4155,7 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { }); if filtered != Type::Union(union) { let place = self.expect_place(&subscript_place_expr); - constraints.insert(place, NarrowingConstraint::replacement(filtered)); + constraints.insert(place, NarrowingConstraint::precomputed_refinement(filtered)); } } @@ -3908,7 +4207,7 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { insert_narrowing_constraint( &mut constraints, self.expect_place(&target), - NarrowingConstraint::replacement(narrowed), + NarrowingConstraint::precomputed_refinement(narrowed), ); } }; @@ -3996,9 +4295,7 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { } } - // Narrow types when a key membership test proves that a key is present, and narrow unions - // and intersections of `TypedDict` when a key membership test proves that a required key is - // absent: + // Record literal key-membership facts for `TypedDict` and `Mapping` values: // // class Foo(TypedDict): // foo: int @@ -4012,81 +4309,26 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { && let Some(key) = inference.expression_type(&**left).as_string_literal() && let rhs_expr = comparators[0].expression_value() && let rhs_type = inference.expression_type(&comparators[0]) - && is_or_contains_typeddict(db, rhs_type) + && is_or_contains_mapping(db, &self.env, rhs_type) { - let key = key.value(db); - let apply_constraint = - |constraints: &mut NarrowingConstraints<'db>, - constraint: NarrowingConstraint<'db>| { - let comparator_place = PlaceExpr::try_from_expr(&comparators[0]) - .and_then(|place_expr| self.places().place_id(&place_expr)); - if let Some(place) = comparator_place { - constraints.insert(place, constraint.clone()); - } - - let value_place = PlaceExpr::try_from_expr(rhs_expr) - .and_then(|place_expr| self.places().place_id(&place_expr)); - if value_place != comparator_place - && let Some(place) = value_place - { - constraints.insert(place, constraint); - } - }; - - if is_positive == (ops[0] == ast::CmpOp::In) { - let narrowed = self.narrow_with_present_key(rhs_type, key); - if narrowed != rhs_type.resolve_type_alias(db) { - apply_constraint(&mut constraints, NarrowingConstraint::replacement(narrowed)); - } + let constraint = if is_positive == (ops[0] == ast::CmpOp::In) { + NarrowingConstraint::present_key(rhs_type, key) } else { - let requires_key = |td: TypedDictType<'db>| -> bool { - td.items(db) - .get(key) - .is_some_and(TypedDictField::is_required) - }; - - let resolved_rhs_type = rhs_type.resolve_type_alias(db); + NarrowingConstraint::absent_key(rhs_type, key) + }; - let narrowed = match resolved_rhs_type { - Type::TypedDict(td) => { - if requires_key(td) { - Type::Never - } else { - resolved_rhs_type - } - } - Type::Intersection(intersection) => { - if intersection - .positive(db) - .iter() - .copied() - .filter_map(Type::as_typed_dict) - .any(requires_key) - { - Type::Never - } else { - resolved_rhs_type - } - } - Type::Union(union) => { - // remove all members of the union that would require the key - union.filter(db, |ty| match ty { - Type::TypedDict(td) => !requires_key(*td), - Type::Intersection(intersection) => !intersection - .positive(db) - .iter() - .copied() - .filter_map(Type::as_typed_dict) - .any(requires_key), - _ => true, - }) - } - _ => resolved_rhs_type, - }; + let comparator_place = PlaceExpr::try_from_expr(&comparators[0]) + .and_then(|place_expr| self.places().place_id(&place_expr)); + if let Some(place) = comparator_place { + constraints.insert(place, constraint.clone()); + } - if narrowed != resolved_rhs_type { - apply_constraint(&mut constraints, NarrowingConstraint::replacement(narrowed)); - } + let value_place = PlaceExpr::try_from_expr(rhs_expr) + .and_then(|place_expr| self.places().place_id(&place_expr)); + if value_place != comparator_place + && let Some(place) = value_place + { + constraints.insert(place, constraint); } } @@ -4258,9 +4500,9 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { let [first_arg, second_arg] = &*expr_call.arguments.args else { return None; }; - let first_arg = PlaceExpr::try_from_expr(first_arg)?; + let first_arg_place = PlaceExpr::try_from_expr(first_arg)?; let function = function_type.known(db)?; - let place = self.expect_place(&first_arg); + let place = self.expect_place(&first_arg_place); if function == KnownFunction::HasAttr { let attr = inference @@ -4272,6 +4514,42 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { return None; } + let places = self.places(); + let receiver = inference.expression_type(first_arg); + if !is_positive + && let Some(member) = places.member_id_by_instance_attribute_name(attr) + && places.place(member).is_bound() + && places + .parents(places.place(member)) + .any(|parent| parent == place) + && matches!(receiver, Type::TypeVar(typevar) if typevar.typevar(db).is_self(db)) + && matches!( + self.predicate, + PredicateNode::Expression(expression) + if !matches!( + expression.node_ref(db).node(self.module), + ast::Expr::BoolOp(_) + ) + ) + && self.predicate_initializes_member(member.into()) + && !receiver.nominal_class(db, &self.env).is_some_and(|class| { + class + .iter_mro(db) + .filter_map(ClassBase::into_class) + .filter_map(|base| base.static_class_literal(db)) + .any(|(base, _)| { + let places = place_table(db, base.body_scope(db)); + places + .symbol_id(attr) + .is_some_and(|symbol| places.symbol(symbol).is_bound()) + }) + }) + { + // An implicit attribute cannot establish its own absence while its + // initializer is being inferred. + return None; + } + // Since `hasattr` only checks if an attribute is readable, // the type of the protocol member should be a read-only property that returns `object`. let constraint = Type::protocol_with_readonly_members( @@ -4366,7 +4644,7 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { let (_, place) = type_guard.place_info(db)?; Some(( place, - NarrowingConstraint::replacement(type_guard.return_type(db)), + NarrowingConstraint::type_guard_replacement(type_guard.return_type(db)), )) } _ => None, @@ -4800,30 +5078,6 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { Some((place, NarrowingConstraint::intersection(intersection))) } - // TODO: Restructure this helper to return the key-presence constraint and apply it with - // `NarrowingConstraint::intersection` at the call site instead of constructing a replacement - // type here. - fn narrow_with_present_key(&self, ty: Type<'db>, key: &str) -> Type<'db> { - let db = self.db; - let constrain = |ty, key_presence_constraint| { - IntersectionType::from_two_elements(db, &self.env, ty, key_presence_constraint) - }; - - match ty.resolve_type_alias(db) { - Type::Union(union) => union.map(db, &self.env, |element| { - self.narrow_with_present_key(*element, key) - }), - resolved if typeddict_declares_key(db, resolved, key) => resolved, - // TODO: Extend this to subtypes of `Mapping[str, object]` whose membership and - // subscript operations obey the `Mapping` contract. - resolved if is_or_contains_typeddict(db, resolved) => constrain( - ty, - Type::TypedDict(required_typeddict_key(db, key, Type::object())), - ), - _ => constrain(ty, key_membership_contains_protocol(db, &self.env, key)), - } - } - /// Narrow tagged unions of tuples with `Literal` elements. /// /// Given a subscript expression like `t[0]` where `t` is a union of tuple types, and a @@ -4894,7 +5148,7 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { // Only create a constraint if we actually narrowed something. if filtered != Type::Union(union) { let place = self.expect_place(&subscript_place_expr); - Some((place, NarrowingConstraint::replacement(filtered))) + Some((place, NarrowingConstraint::precomputed_refinement(filtered))) } else { None } @@ -4946,7 +5200,7 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { let attribute_value_place_expr = PlaceExpr::try_from_expr(attribute_value_expr)?; let place = self.expect_place(&attribute_value_place_expr); - Some((place, NarrowingConstraint::replacement(narrowed))) + Some((place, NarrowingConstraint::precomputed_refinement(narrowed))) } fn narrow_nominal_attribute_by_truthiness( @@ -4983,7 +5237,7 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { let attribute_value_place_expr = PlaceExpr::try_from_expr(attribute_value_expr)?; let place = self.expect_place(&attribute_value_place_expr); - Some((place, NarrowingConstraint::replacement(narrowed))) + Some((place, NarrowingConstraint::precomputed_refinement(narrowed))) } } @@ -5034,58 +5288,194 @@ fn is_or_contains_typeddict<'db>(db: &'db dyn Db, ty: Type<'db>) -> bool { } } -fn typeddict_declares_key<'db>(db: &'db dyn Db, ty: Type<'db>, key: &str) -> bool { - match ty { - Type::TypedDict(typed_dict) => typed_dict.items(db).contains_key(key), - Type::Intersection(intersection) => intersection - .positive(db) +fn is_or_contains_mapping<'db>( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + ty: Type<'db>, +) -> bool { + match ty.resolve_type_alias(db) { + Type::Union(union) => union + .elements(db) .iter() - .any(|element| typeddict_declares_key(db, *element, key)), + .any(|element| is_or_contains_mapping(db, env, *element)), + resolved => is_mapping_subtype(db, env, resolved), + } +} + +fn is_mapping_subtype<'db>(db: &'db dyn Db, env: &ProgramEnvironment<'db>, ty: Type<'db>) -> bool { + ty.is_subtype_of( + db, + env, + KnownClass::Mapping + .to_instance(db, env) + .top_materialization(db, env), + ) +} + +/// Return whether successful membership implies that subscripting with the same key is valid. +/// +/// `Mapping` provides this relationship; an arbitrary `__contains__` implementation does not. +/// Every arm of a union must provide the relationship before it can be carried across replacement +/// narrowing. +fn key_membership_implies_subscript<'db>( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + ty: Type<'db>, +) -> bool { + match ty.resolve_type_alias(db) { Type::Union(union) => union .elements(db) .iter() - .any(|element| typeddict_declares_key(db, *element, key)), - Type::TypeAlias(alias) => typeddict_declares_key(db, alias.value_type(db), key), + .all(|element| key_membership_implies_subscript(db, env, *element)), + resolved => is_mapping_subtype(db, env, resolved), + } +} + +/// Refine `ty` with the fact that the literal `key` is present. +/// +/// `Mapping` arms gain matching `__contains__` and `__getitem__` methods, while other arms only +/// gain the successful membership fact. +fn narrow_with_present_key<'db>( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + ty: Type<'db>, + key: &str, +) -> Type<'db> { + let constrain = |ty, key_presence_constraint| { + IntersectionType::from_two_elements(db, env, ty, key_presence_constraint) + }; + + match ty.resolve_type_alias(db) { + Type::Union(union) => union.map(db, env, |element| { + narrow_with_present_key(db, env, *element, key) + }), + resolved if closed_typeddict_excludes_key(db, resolved, key) => Type::Never, + resolved if is_mapping_subtype(db, env, resolved) => { + constrain(ty, mapping_present_key_protocol(db, env, key)) + } + _ => constrain(ty, key_membership_contains_protocol(db, env, key)), + } +} + +/// Return whether a closed `TypedDict` in `ty` rules out `key`. +fn closed_typeddict_excludes_key<'db>(db: &'db dyn Db, ty: Type<'db>, key: &str) -> bool { + match ty.resolve_type_alias(db) { + Type::TypedDict(typed_dict) => { + typed_dict.openness(db).is_closed() && !typed_dict.items(db).contains_key(key) + } + Type::Intersection(intersection) => intersection + .positive(db) + .iter() + .any(|element| closed_typeddict_excludes_key(db, *element, key)), _ => false, } } -/// Return a synthesized `TypedDict` that represents safe subscript access for a present key on a -/// `TypedDict`-containing type. +/// Refine `ty` with the fact that the literal `key` is absent. /// -/// For `TypedDict`s, a positive key-membership test proves more than containment: it also makes -/// string-literal subscript access with that key valid. In the `if` branch below, the `Bar` arm -/// keeps its original shape but is intersected with this schema so `u["foo"]` is accepted: +/// Union arms that prove the key is always present are removed. Other arms are retained unchanged: +/// failing to prove presence is not proof of absence. +fn narrow_with_absent_key<'db>( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + ty: Type<'db>, + key: StringLiteralType<'db>, +) -> Type<'db> { + let resolved = ty.resolve_type_alias(db); + match resolved { + Type::Union(union) => { + let filtered = union.filter(db, |ty| !key_is_always_present(db, env, *ty, key)); + if filtered == resolved { ty } else { filtered } + } + ty if key_is_always_present(db, env, ty, key) => Type::Never, + _ => ty, + } +} + +/// Return whether `ty` proves that `key` is present, making a negative membership branch +/// unreachable. /// -/// ```python -/// class Foo(TypedDict): -/// foo: int +/// Calling `__contains__` with this literal key must have an always-truthy return type. +fn key_is_always_present<'db>( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + ty: Type<'db>, + key: StringLiteralType<'db>, +) -> bool { + let resolved = ty.resolve_type_alias(db); + if let Type::Intersection(intersection) = resolved + && intersection + .positive(db) + .iter() + .any(|element| key_is_always_present(db, env, *element, key)) + { + return true; + } + + resolved + .try_call_dunder( + db, + env, + "__contains__", + CallArguments::positional([Type::string_literal(db, key.value(db))]), + TypeContext::default(), + ) + .is_ok_and(|bindings| bindings.return_type(db, env).bool(db, env) == Truthiness::AlwaysTrue) +} + +/// Return a synthesized protocol that records a present key on a `Mapping` subtype. /// -/// class Bar(TypedDict): -/// bar: int +/// This preserves both the proven membership result and safe subscript access, assuming that +/// `Mapping` implementations honor the contract between the two operations. /// -/// def f(u: Foo | Bar): -/// if "foo" in u: -/// reveal_type(u["foo"]) # object +/// ```python +/// def f(mapping: Mapping[Literal["a"], int]): +/// if "b" in mapping: +/// reveal_type(mapping["b"]) # object /// ``` -fn required_typeddict_key<'db>( +fn mapping_present_key_protocol<'db>( db: &'db dyn Db, + env: &ProgramEnvironment<'db>, key: &str, - value_ty: Type<'db>, -) -> TypedDictType<'db> { - let field = TypedDictFieldBuilder::new(value_ty) - .required(true) - .read_only(true) - .build(); - let schema = TypedDictSchema::from_iter([(Name::from(key), field)]); - TypedDictType::from_schema_items(db, schema) +) -> Type<'db> { + let getitem_signature = Signature::new( + Parameters::standard([ + Parameter::positional_only(Some(Name::new_static("self"))), + Parameter::positional_only(Some(Name::new_static("key"))) + .with_annotated_type(Type::string_literal(db, key)), + ]), + Type::object(), + ); + let contains_signature = Signature::new( + Parameters::standard([ + Parameter::positional_only(Some(Name::new_static("self"))), + Parameter::positional_only(Some(Name::new_static("key"))) + .with_annotated_type(Type::string_literal(db, key)), + ]), + Type::bool_literal(true), + ); + + Type::protocol_with_methods( + db, + env, + [ + ( + "__getitem__", + CallableType::function_like(db, getitem_signature), + ), + ( + "__contains__", + CallableType::function_like(db, contains_signature), + ), + ], + ) } /// Return a synthesized protocol that records a true key-membership test without implying /// subscript access. /// -/// For non-`TypedDict` types, `"key" in value` only proves that membership is true. It does not -/// prove that `value["key"]` is valid: +/// For non-`Mapping` types, `"key" in value` only proves that membership is true. It does not prove +/// that `value["key"]` is valid: /// /// ```python /// def f(s: Literal["abc"]): @@ -5093,8 +5483,8 @@ fn required_typeddict_key<'db>( /// s["a"] # Runtime `TypeError` /// ``` /// -/// Non-`TypedDict` union arms therefore receive this `__contains__` protocol instead of the -/// synthesized `TypedDict` used for `TypedDict` arms. +/// Non-`Mapping` union arms therefore receive this `__contains__` protocol instead of the +/// `__contains__` and `__getitem__` protocol used for `Mapping` arms. fn key_membership_contains_protocol<'db>( db: &'db dyn Db, env: &ProgramEnvironment<'db>,