diff --git a/crates/ty_python_semantic/resources/mdtest/comparison/instances/membership_test.md b/crates/ty_python_semantic/resources/mdtest/comparison/instances/membership_test.md index d101b57fdebcf..865d2e55e17a2 100644 --- a/crates/ty_python_semantic/resources/mdtest/comparison/instances/membership_test.md +++ b/crates/ty_python_semantic/resources/mdtest/comparison/instances/membership_test.md @@ -10,6 +10,172 @@ For references, see: - - +## Inline list and set literals + +An immediately consumed list or set literal preserves its elements' values for membership tests. +This lets us determine whether a literal value is present or absent. + +```py +reveal_type("a" in ["a", "b"]) # revealed: Literal[True] +reveal_type("c" in ["a", "b"]) # revealed: Literal[False] +reveal_type("a" not in ["a", "b"]) # revealed: Literal[False] +reveal_type("c" not in ["a", "b"]) # revealed: Literal[True] + +reveal_type("a" in {"a", "b"}) # revealed: Literal[True] +reveal_type("c" in {"a", "b"}) # revealed: Literal[False] +reveal_type("a" not in {"a", "b"}) # revealed: Literal[False] +reveal_type("c" not in {"a", "b"}) # revealed: Literal[True] +``` + +Membership uses runtime equality, including equality between integer and boolean values. Duplicate +set elements do not change membership, and an empty list cannot contain any value. + +```py +reveal_type(1 in [True]) # revealed: Literal[True] +reveal_type(True in {1, 1}) # revealed: Literal[True] +reveal_type(b"a" in {b"a", b"b"}) # revealed: Literal[True] +reveal_type("a" in []) # revealed: Literal[False] +reveal_type("a" not in []) # revealed: Literal[True] +``` + +## Inline elements with inferred literal types + +Elements can have known literal values without using literal syntax. A function returning a literal +type provides the same membership information as the corresponding literal expression. + +```py +from typing import Literal + +def first() -> Literal["a"]: + return "a" + +second = "b" +reveal_type("a" in [first(), second]) # revealed: Literal[True] +reveal_type("c" in [first(), second]) # revealed: Literal[False] +reveal_type("b" in {first(), second}) # revealed: Literal[True] +reveal_type("c" in {first(), second}) # revealed: Literal[False] +``` + +An element with a union type contributes only one of its possible values. Neither alternative is +guaranteed to be present, but a value outside the union is definitely absent. + +```py +def uncertain_element(value: Literal["a", "b"]): + reveal_type("a" in [value]) # revealed: bool + reveal_type("c" in [value]) # revealed: Literal[False] + reveal_type("a" in {value}) # revealed: bool + reveal_type("c" in {value}) # revealed: Literal[False] + +def unknown_element(value: str): + reveal_type("a" in [value]) # revealed: bool + reveal_type("a" in {value}) # revealed: bool +``` + +An unknown element does not prevent a later known element from establishing membership. + +```py +def known_element(value: str): + reveal_type("a" in [value, "a"]) # revealed: Literal[True] + reveal_type("a" not in [value, "a"]) # revealed: Literal[False] + reveal_type("a" in {value, "a"}) # revealed: Literal[True] + reveal_type("a" not in {value, "a"}) # revealed: Literal[False] +``` + +## Inline enum members + +An inline set of enum members also has known membership. + +```py +from enum import Enum + +class Color(Enum): + RED = 1 + BLUE = 2 + +reveal_type(Color.RED in {Color.RED}) # revealed: Literal[True] +reveal_type(Color.BLUE in {Color.RED}) # revealed: Literal[False] +``` + +## Unpacked inline list elements + +Unpacking a fixed-length tuple or another inline list preserves the known elements. Unpacking a list +of unknown length does not guarantee that any value is present. + +```py +reveal_type("a" in [*("a", "b")]) # revealed: Literal[True] +reveal_type("c" not in [*["a", "b"]]) # revealed: Literal[True] + +def unpack_unknown(values: list[str]): + reveal_type("a" in [*values]) # revealed: bool +``` + +## Stored mutable containers + +A stored list or set can be empty or contain only some of the values permitted by its element type. +A literal element type alone therefore does not establish membership. + +```py +from typing import Literal + +def stored(values: list[Literal["a", "b"]], keys: set[Literal["a", "b"]]): + reveal_type("a" in values) # revealed: bool + reveal_type("a" not in values) # revealed: bool + reveal_type("a" in keys) # revealed: bool + reveal_type("a" not in keys) # revealed: bool +``` + +## Inline containers with custom equality + +Lists and sets test their elements using identity or equality. A custom equality method with a known +result can therefore establish membership. We assume that equal objects have equal hashes. + +```py +from typing import Literal, overload + +class EqualValue: + @overload + def __eq__(self, other: "EqualValue") -> Literal[True]: ... + @overload + def __eq__(self, other: object) -> bool: ... + def __eq__(self, other: object) -> bool: + return isinstance(other, EqualValue) + + def __hash__(self) -> int: + return 0 + +reveal_type(EqualValue() in [EqualValue()]) # revealed: Literal[True] +reveal_type(EqualValue() not in [EqualValue()]) # revealed: Literal[False] +reveal_type(EqualValue() in {EqualValue()}) # revealed: Literal[True] +reveal_type(EqualValue() not in {EqualValue()}) # revealed: Literal[False] +``` + +Identity also counts as a match, even when equality always returns false. The equality result alone +therefore cannot establish that a custom object is absent. + +```py +class NeverEqual: + def __eq__(self, other: object) -> Literal[False]: + return False + +value = NeverEqual() +reveal_type(value in [value]) # revealed: bool +``` + +## Inline membership in comparison chains + +The following comparison receives the actual list type of the shared operand, so its reflected +method can accept the list even when the preceding membership test has a known result. + +```py +from typing import Literal + +class LargerThanList: + def __gt__(self, other: list[str]) -> Literal[True]: + return True + +reveal_type("a" in ["a"] < LargerThanList()) # revealed: Literal[True] +``` + ## Implements `__contains__` Classes can support membership tests by implementing the `__contains__` method: diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/in.md b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/in.md index af0c5666356cb..6eb4e17a924e9 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/in.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/in.md @@ -1368,6 +1368,21 @@ def after_excluding_red_mixed(x: Color | int): reveal_type(x) # revealed: Literal[Color.BLUE] | int ``` +Inline set literals also preserve the individual enum members for narrowing in either branch. + +```py +def inline_enum_set(x: Color): + if x in {Color.RED, Color.GREEN}: + reveal_type(x) # revealed: Literal[Color.RED, Color.GREEN] + else: + reveal_type(x) # revealed: Literal[Color.BLUE] + + if x not in {Color.RED}: + reveal_type(x) # revealed: Literal[Color.GREEN, Color.BLUE] + else: + reveal_type(x) # revealed: Literal[Color.RED] +``` + When the container's element type is a union of enum literals, membership narrows to that union. Without the annotation, the tuple's elements are widened to `Color`, so the comprehension remains `list[Color]`: diff --git a/crates/ty_python_semantic/resources/mdtest/statically_known_branches.md b/crates/ty_python_semantic/resources/mdtest/statically_known_branches.md index 4f5cc43ed2c49..3a56155da461d 100644 --- a/crates/ty_python_semantic/resources/mdtest/statically_known_branches.md +++ b/crates/ty_python_semantic/resources/mdtest/statically_known_branches.md @@ -106,6 +106,50 @@ darwin other ``` +### `sys.platform` membership in inline containers + +Membership in an inline list or set of platform names is statically known when the target platform +is configured. Assignments in matching branches replace the initial value. + +```toml +[environment] +python-platform = "linux" +``` + +```py +import sys + +x = 0 +if sys.platform in {"linux", "windows"}: + x = 1 +reveal_type(x) # revealed: Literal[1] + +y = 0 +if sys.platform in ["linux", "windows"]: + y = 1 +reveal_type(y) # revealed: Literal[1] +``` + +An absent platform leaves the initial value unchanged. Negating a membership test reverses which +branch is reachable. + +```py +absent = 0 +if sys.platform in ["darwin", "win32"]: + absent = 1 +reveal_type(absent) # revealed: Literal[0] + +excluded = 0 +if sys.platform not in {"darwin", "win32"}: + excluded = 1 +reveal_type(excluded) # revealed: Literal[1] + +present = 0 +if sys.platform not in ["linux", "win32"]: + present = 1 +reveal_type(present) # revealed: Literal[0] +``` + ### `typing.TYPE_CHECKING` ```py diff --git a/crates/ty_python_semantic/src/types/equality.rs b/crates/ty_python_semantic/src/types/equality.rs index 91d82d553d69c..aa6a50bde23c8 100644 --- a/crates/ty_python_semantic/src/types/equality.rs +++ b/crates/ty_python_semantic/src/types/equality.rs @@ -296,14 +296,15 @@ pub(super) fn inequality_truthiness<'db>( ) } -/// Evaluates tuple-element equality while reusing the active-comparison-set allocation across a -/// tuple walk. The set only detects recursive comparisons; results are not cached between -/// elements. -pub(super) struct TupleEqualityEvaluator<'db> { +/// Evaluates container elements using identity or equality. +/// +/// Reuses the active-comparison set across element comparisons. The set only detects recursive +/// comparisons; results are not cached between elements. +pub(super) struct ContainerElementEqualityEvaluator<'db> { evaluator: ComparisonEvaluator<'db>, } -impl<'db> TupleEqualityEvaluator<'db> { +impl<'db> ContainerElementEqualityEvaluator<'db> { pub(super) fn new( db: &'db dyn Db, env: &ProgramEnvironment<'db>, @@ -320,7 +321,7 @@ impl<'db> TupleEqualityEvaluator<'db> { right: Type<'db>, ) -> Result> { let db = self.evaluator.db; - let truthiness = evaluate_tuple_element_equality(&mut self.evaluator, left, right); + let truthiness = evaluate_container_element_equality(&mut self.evaluator, left, right); if !truthiness.is_ambiguous() { return Ok(truthiness); } @@ -1671,7 +1672,7 @@ fn compare_nominal_instances<'db>( let mut all_equal = true; for (&left, &right) in left_elements.iter().zip(right_elements) { - match evaluate_tuple_element_equality(evaluator, left, right) { + match evaluate_container_element_equality(evaluator, left, right) { Truthiness::AlwaysTrue => {} Truthiness::AlwaysFalse => return operator.result_from_equality(false), Truthiness::Ambiguous => all_equal = false, @@ -1688,7 +1689,7 @@ fn compare_nominal_instances<'db>( } } -fn evaluate_tuple_element_equality<'db>( +fn evaluate_container_element_equality<'db>( evaluator: &mut ComparisonEvaluator<'db>, left: Type<'db>, right: Type<'db>, @@ -1928,6 +1929,16 @@ fn has_reflexive_equality_semantics<'db>( evaluator: &ComparisonEvaluator<'db>, ty: Type<'db>, ) -> bool { + // `comparison_semantics` identifies a single builtin comparison implementation. The evaluator + // normally distributes unions before that lookup, but this check receives the original operands. + // A union such as `Literal[1, "a"]` has reflexive equality in every alternative without sharing + // one comparison implementation, so check each alternative here. + if let Type::Union(union) = ty.resolve_type_alias(evaluator.db) { + return union + .elements(evaluator.db) + .iter() + .all(|element| has_reflexive_equality_semantics(evaluator, *element)); + } evaluator .comparison_semantics(ty, ComparisonOperator::Equality) .is_some() diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 842be8d63926d..b309329d806d6 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -11644,14 +11644,25 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let range = TextRange::new(left.start(), right.end()); - let ty = comparisons::infer_binary_type_comparison( + let literal_membership = comparisons::infer_literal_membership_comparison( &builder.context, left_ty, *op, - right_ty, - range, - ) - .unwrap_or_else(|error| { + right, + |element| builder.expression_type(element), + ); + + let comparison = match literal_membership { + Some(ty) => Ok(ty), + None => comparisons::infer_binary_type_comparison( + &builder.context, + left_ty, + *op, + right_ty, + range, + ), + }; + let ty = comparison.unwrap_or_else(|error| { report_unsupported_comparison( &builder.context, &error, diff --git a/crates/ty_python_semantic/src/types/infer/comparisons.rs b/crates/ty_python_semantic/src/types/infer/comparisons.rs index cdfa41a824bf9..d9fbc295f206c 100644 --- a/crates/ty_python_semantic/src/types/infer/comparisons.rs +++ b/crates/ty_python_semantic/src/types/infer/comparisons.rs @@ -11,8 +11,10 @@ use crate::types::constraints::ConstraintSetBuilder; use crate::types::context::InferContext; use crate::types::cyclic::CycleDetector; use crate::types::equality::{ - ComparisonSoundnessPolicy, TupleEqualityEvaluator, equality_truthiness, inequality_truthiness, + ComparisonSoundnessPolicy, ContainerElementEqualityEvaluator, equality_truthiness, + inequality_truthiness, }; +use crate::types::iteration::extract_literal_container_element_types; use crate::types::known_instance::{FunctoolsPartialInstance, InternedType, MethodWrapper}; use crate::types::tuple::{Tuple, TupleSpec}; use crate::types::{ @@ -646,10 +648,6 @@ enum MembershipOperator { } impl MembershipOperator { - const fn is_in(self) -> bool { - matches!(self, MembershipOperator::In) - } - const fn is_not_in(self) -> bool { matches!(self, MembershipOperator::NotIn) } @@ -697,6 +695,54 @@ pub(crate) struct UnsupportedComparisonError<'db> { pub(crate) right_ty: Type<'db>, } +/// Refine membership using the contents of an immediately consumed container display. +/// For sets, assume equality is an equivalence relation and equal objects have equal hashes. +pub(super) fn infer_literal_membership_comparison<'db>( + context: &InferContext<'db, '_>, + left: Type<'db>, + op: ast::CmpOp, + right: &ast::Expr, + expression_type: impl FnMut(&ast::Expr) -> Type<'db>, +) -> Option> { + let negate = match op { + ast::CmpOp::In => false, + ast::CmpOp::NotIn => true, + _ => return None, + }; + let db = context.db(); + let env = context.program_environment(); + let elements = extract_literal_container_element_types(db, env, right, expression_type)?; + let truthiness = fixed_membership_truthiness(context, left, &elements).negate_if(negate); + Some(Type::from_truthiness(db, env, truthiness)) +} + +/// Evaluate membership using the supplied element types and identity-or-equality semantics. +/// An ambiguous comparison does not prevent a later element from proving membership. +fn fixed_membership_truthiness<'db>( + context: &InferContext<'db, '_>, + needle: Type<'db>, + elements: &[Type<'db>], +) -> Truthiness { + let db = context.db(); + let env = context.program_environment(); + let soundness_policy = + ComparisonSoundnessPolicy::from_analysis_settings(db.analysis_settings(context.file())); + let mut equality = ContainerElementEqualityEvaluator::new(db, env, soundness_policy); + let mut truthiness = Truthiness::AlwaysFalse; + for &element in elements { + // It's okay to ignore errors here because Python doesn't call `__bool__` + // for different union variants. Instead, this is just for us to + // evaluate a possibly truthy value to `false` or `true`. + truthiness = truthiness.or(equality + .element_truthiness(element, needle) + .unwrap_or_else(|error| error.fallback_truthiness())); + if truthiness.is_always_true() { + break; + } + } + truthiness +} + /// Infers the type of a binary comparison (e.g. 'left == right'). See /// `TypeInferenceBuilder::infer_compare_expression` for the higher level logic dealing with /// multi-comparison expressions. @@ -781,31 +827,9 @@ fn infer_binary_type_comparison_inner<'db>( && let Some(right_tuple) = right.tuple_instance_spec(db, env) && let Tuple::Fixed(right_tuple) = &*right_tuple { - let mut any_eq = false; - let mut any_ambiguous = false; - let mut equality = TupleEqualityEvaluator::new(db, env, soundness_policy); - - for &element_ty in right_tuple.elements_slice() { - // It's okay to ignore errors here because Python doesn't call `__bool__` - // for different union variants. Instead, this is just for us to - // evaluate a possibly truthy value to `false` or `true`. - match equality - .element_truthiness(element_ty, left) - .unwrap_or_else(|error| error.fallback_truthiness()) - { - Truthiness::AlwaysTrue => any_eq = true, - Truthiness::AlwaysFalse => (), - Truthiness::Ambiguous => any_ambiguous = true, - } - } - - return Ok(if any_eq { - Type::bool_literal(op.is_in()) - } else if !any_ambiguous { - Type::bool_literal(op.is_not_in()) - } else { - KnownClass::Bool.to_instance(db, env) - }); + let truthiness = fixed_membership_truthiness(context, left, right_tuple.elements_slice()) + .negate_if(op.is_not_in()); + return Ok(Type::from_truthiness(db, env, truthiness)); } let comparison_truthiness = match op { @@ -1552,7 +1576,7 @@ fn infer_tuple_rich_comparison<'db>( let soundness_policy = ComparisonSoundnessPolicy::from_analysis_settings( db.analysis_settings(context.file()), ); - let mut equality = TupleEqualityEvaluator::new(db, env, soundness_policy); + let mut equality = ContainerElementEqualityEvaluator::new(db, env, soundness_policy); for (l_ty, r_ty) in left_iter.zip(right_iter) { let eq_truthiness = equality diff --git a/crates/ty_python_semantic/src/types/iteration.rs b/crates/ty_python_semantic/src/types/iteration.rs index faecaa25a636b..cb96e8851034c 100644 --- a/crates/ty_python_semantic/src/types/iteration.rs +++ b/crates/ty_python_semantic/src/types/iteration.rs @@ -15,6 +15,27 @@ use ruff_python_ast as ast; use std::borrow::Cow; use ty_python_core::EvaluationMode; +/// Extract precise element types for membership in an immediately consumed list or set display. +/// +/// The ordinary inferred container type remains unchanged. Set construction can remove duplicates, +/// so these elements do not describe an iteration order or an exact length. +pub(super) fn extract_literal_container_element_types<'db>( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + expression: &ast::Expr, + expression_type: impl FnMut(&ast::Expr) -> Type<'db>, +) -> Option]>> { + match expression { + ast::Expr::List(_) => { + extract_fixed_length_iterable_element_types(db, env, expression, expression_type) + } + ast::Expr::Set(set) if !set.elts.iter().any(ast::Expr::is_starred_expr) => { + Some(set.elts.iter().map(expression_type).collect()) + } + _ => None, + } +} + /// Extract the element types from an expression with a statically known fixed-length iteration. /// /// List and tuple literals are expanded directly so we preserve precise element types, including diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index 6bbd96df3ccab..304fa7db02ac2 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -9,6 +9,7 @@ use crate::reachability::{ use crate::subscript::PyIndex; use crate::types::function::KnownFunction; use crate::types::infer::{ExpressionInference, infer_same_file_expression_type}; +use crate::types::iteration::extract_literal_container_element_types; use crate::types::special_form::TypeQualifier; use crate::types::tuple::{TupleElement, TupleLength, TupleSpec, TupleSpecBuilder, TupleType}; use crate::types::typed_dict::{TypedDictFieldBuilder, TypedDictSchema, TypedDictType}; @@ -3786,30 +3787,21 @@ impl<'db> NarrowingConstraintsBuilder<'db, '_> { /// Preserve the precise element types of an immediately consumed list or set literal. /// - /// These expressions cannot be mutated before membership is evaluated. Representing them as - /// fixed-length tuples also lets negative narrowing exclude values that are guaranteed present. + /// A tuple carries the element types into equality-based narrowing, including negative + /// constraints for singleton elements. This view does not describe a set's length or order. fn inline_membership_rhs_type( &self, rhs: &ast::Expr, inference: &ExpressionInference<'db>, ) -> Option> { let db = self.db; - let elements = match rhs.expression_value() { - ast::Expr::List(list) => &list.elts, - ast::Expr::Set(set) => &set.elts, - _ => return None, - }; - - if elements.iter().any(ast::Expr::is_starred_expr) { - return None; - } - + let elements = extract_literal_container_element_types(db, &self.env, rhs, |element| { + inference.expression_type(element) + })?; Some(Type::heterogeneous_tuple( db, &self.env, - elements - .iter() - .map(|element| inference.expression_type(element)), + elements.iter().copied(), )) }