From 9212e03c4b008bad81ff8923bb5964683d634321 Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Sat, 26 Apr 2025 18:27:10 +0900 Subject: [PATCH 01/17] [red-knot] basic narrowing on attribute and subscript expressions --- .../mdtest/narrow/composite_target.md | 73 +++++++ .../src/semantic_index/definition.rs | 97 ++++++--- .../src/semantic_index/use_def.rs | 6 +- crates/ty_python_semantic/src/symbol.rs | 2 +- crates/ty_python_semantic/src/types/infer.rs | 56 ++++- crates/ty_python_semantic/src/types/narrow.rs | 191 ++++++++++++------ 6 files changed, 320 insertions(+), 105 deletions(-) create mode 100644 crates/ty_python_semantic/resources/mdtest/narrow/composite_target.md diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/composite_target.md b/crates/ty_python_semantic/resources/mdtest/narrow/composite_target.md new file mode 100644 index 00000000000000..140787dcf19dd8 --- /dev/null +++ b/crates/ty_python_semantic/resources/mdtest/narrow/composite_target.md @@ -0,0 +1,73 @@ +# Narrowing for composite targets + +## Member access + +```py +class C: + def __init__(self): + self.x: int | None = None + self.y: int | None = None + +c = C() +reveal_type(c.x) # revealed: int | None +if c.x is not None: + reveal_type(c.x) # revealed: int + reveal_type(c.y) # revealed: int | None + +if c.x is not None: + def _(): + reveal_type(c.x) # revealed: Unknown | int | None + +def _(): + if c.x is not None: + reveal_type(c.x) # revealed: (Unknown & ~None) | int +``` + +## Subscript + +### Number subscript + +```py +def _(t1: tuple[int | None, int | None], t2: tuple[int, int] | tuple[None, None]): + if t1[0] is not None: + reveal_type(t1[0]) # revealed: int + reveal_type(t1[1]) # revealed: int | None + + n = 0 + if t1[n] is not None: + # Non-literal subscript narrowing are currently not supported, as well as mypy, pyright + reveal_type(t1[0]) # revealed: int | None + reveal_type(t1[n]) # revealed: int | None + reveal_type(t1[1]) # revealed: int | None + + if t2[0] is not None: + # TODO: should be int + reveal_type(t2[0]) # revealed: Unknown & ~None + # TODO: should be int + reveal_type(t2[1]) # revealed: Unknown +``` + +### String subscript + +```py +def _(d: dict[str, str | None]): + if d["a"] is not None: + reveal_type(d["a"]) # revealed: str + reveal_type(d["b"]) # revealed: str | None +``` + +## Complex target + +```py +class C: + def __init__(self): + self.x: tuple[int | None, int | None] = () + +class D: + def __init__(self): + self.c: tuple[C] | None = None + +d = D() +if d.c is not None and d.c[0].x[0] is not None: + reveal_type(d.c[0].x[0]) # revealed: int +``` diff --git a/crates/ty_python_semantic/src/semantic_index/definition.rs b/crates/ty_python_semantic/src/semantic_index/definition.rs index e0a1a6c571f323..c09afc3fae5020 100644 --- a/crates/ty_python_semantic/src/semantic_index/definition.rs +++ b/crates/ty_python_semantic/src/semantic_index/definition.rs @@ -54,7 +54,7 @@ impl<'db> Definition<'db> { } pub fn focus_range(self, db: &'db dyn Db) -> FileRange { - FileRange::new(self.file(db), self.kind(db).target_range()) + FileRange::new(self.file(db), self.kind(db).target().range()) } } @@ -535,6 +535,25 @@ impl DefinitionCategory { } } +#[derive(Clone, Copy, Debug)] +pub(crate) enum DefinitionTarget<'db> { + Ident(&'db ast::Identifier), + Expr(ast::ExprRef<'db>), +} + +impl DefinitionTarget<'_> { + /// Returns the [`TextRange`] of the definition target. + /// + /// A definition target would mainly be the node representing the symbol being defined i.e., + /// [`ast::ExprName`] or [`ast::Identifier`] but could also be other nodes. + pub(crate) fn range(&self) -> TextRange { + match self { + Self::Ident(ident) => ident.range, + Self::Expr(expr) => expr.range(), + } + } +} + /// The kind of a definition. /// /// ## Usage in salsa tracked structs @@ -583,33 +602,57 @@ impl DefinitionKind<'_> { } } - /// Returns the [`TextRange`] of the definition target. - /// - /// A definition target would mainly be the node representing the symbol being defined i.e., - /// [`ast::ExprName`] or [`ast::Identifier`] but could also be other nodes. - pub(crate) fn target_range(&self) -> TextRange { + pub(crate) fn target(&self) -> DefinitionTarget { match self { - DefinitionKind::Import(import) => import.alias().range(), - DefinitionKind::ImportFrom(import) => import.alias().range(), - DefinitionKind::StarImport(import) => import.alias().range(), - DefinitionKind::Function(function) => function.name.range(), - DefinitionKind::Class(class) => class.name.range(), - DefinitionKind::TypeAlias(type_alias) => type_alias.name.range(), - DefinitionKind::NamedExpression(named) => named.target.range(), - DefinitionKind::Assignment(assignment) => assignment.target.range(), - DefinitionKind::AnnotatedAssignment(assign) => assign.target.range(), - DefinitionKind::AugmentedAssignment(aug_assign) => aug_assign.target.range(), - DefinitionKind::For(for_stmt) => for_stmt.target.range(), - DefinitionKind::Comprehension(comp) => comp.target().range(), - DefinitionKind::VariadicPositionalParameter(parameter) => parameter.name.range(), - DefinitionKind::VariadicKeywordParameter(parameter) => parameter.name.range(), - DefinitionKind::Parameter(parameter) => parameter.parameter.name.range(), - DefinitionKind::WithItem(with_item) => with_item.target.range(), - DefinitionKind::MatchPattern(match_pattern) => match_pattern.identifier.range(), - DefinitionKind::ExceptHandler(handler) => handler.node().range(), - DefinitionKind::TypeVar(type_var) => type_var.name.range(), - DefinitionKind::ParamSpec(param_spec) => param_spec.name.range(), - DefinitionKind::TypeVarTuple(type_var_tuple) => type_var_tuple.name.range(), + DefinitionKind::Import(import) => DefinitionTarget::Ident(&import.alias().name), + DefinitionKind::ImportFrom(import) => DefinitionTarget::Ident(&import.alias().name), + DefinitionKind::StarImport(import) => DefinitionTarget::Ident(&import.alias().name), + DefinitionKind::Function(function) => DefinitionTarget::Ident(&function.name), + DefinitionKind::Class(class) => DefinitionTarget::Ident(&class.name), + DefinitionKind::TypeAlias(type_alias) => { + DefinitionTarget::Expr(ast::ExprRef::from(&type_alias.name)) + } + DefinitionKind::TypeVar(type_var) => DefinitionTarget::Ident(&type_var.name), + DefinitionKind::ParamSpec(param_spec) => DefinitionTarget::Ident(¶m_spec.name), + DefinitionKind::TypeVarTuple(type_var_tuple) => { + DefinitionTarget::Ident(&type_var_tuple.name) + } + DefinitionKind::NamedExpression(named) => { + DefinitionTarget::Expr(ast::ExprRef::from(&named.target)) + } + DefinitionKind::Parameter(parameter) => { + DefinitionTarget::Ident(¶meter.parameter.name) + } + DefinitionKind::VariadicPositionalParameter(parameter) => { + DefinitionTarget::Ident(¶meter.name) + } + DefinitionKind::VariadicKeywordParameter(parameter) => { + DefinitionTarget::Ident(¶meter.name) + } + DefinitionKind::Assignment(assignment) => { + DefinitionTarget::Expr(ast::ExprRef::from(assignment.target.node())) + } + DefinitionKind::AnnotatedAssignment(assignment) => { + DefinitionTarget::Expr(ast::ExprRef::from(assignment.target.node())) + } + DefinitionKind::AugmentedAssignment(aug_assign) => { + DefinitionTarget::Expr(ast::ExprRef::from(&aug_assign.target)) + } + DefinitionKind::For(for_stmt) => { + DefinitionTarget::Expr(ast::ExprRef::from(for_stmt.target.node())) + } + DefinitionKind::Comprehension(comp) => { + DefinitionTarget::Expr(ast::ExprRef::from(comp.target())) + } + DefinitionKind::WithItem(with_item) => { + DefinitionTarget::Expr(ast::ExprRef::from(with_item.target.node())) + } + DefinitionKind::MatchPattern(match_pattern) => { + DefinitionTarget::Ident(&match_pattern.identifier) + } + DefinitionKind::ExceptHandler(handler) => { + DefinitionTarget::Ident(handler.node().name.as_ref().unwrap()) + } } } diff --git a/crates/ty_python_semantic/src/semantic_index/use_def.rs b/crates/ty_python_semantic/src/semantic_index/use_def.rs index c53c8eb4687db3..9af7ca4583ebbb 100644 --- a/crates/ty_python_semantic/src/semantic_index/use_def.rs +++ b/crates/ty_python_semantic/src/semantic_index/use_def.rs @@ -265,7 +265,7 @@ use self::symbol_state::{ }; use crate::node_key::NodeKey; use crate::semantic_index::ast_ids::ScopedUseId; -use crate::semantic_index::definition::Definition; +use crate::semantic_index::definition::{Definition, DefinitionTarget}; use crate::semantic_index::narrowing_constraints::{ ConstraintKey, NarrowingConstraints, NarrowingConstraintsBuilder, NarrowingConstraintsIterator, }; @@ -595,10 +595,10 @@ impl<'db> ConstraintsIterator<'_, 'db> { self, db: &'db dyn crate::Db, base_ty: Type<'db>, - symbol: ScopedSymbolId, + target: DefinitionTarget, ) -> Type<'db> { let constraint_tys: Vec<_> = self - .filter_map(|constraint| infer_narrowing_constraint(db, constraint, symbol)) + .filter_map(|constraint| infer_narrowing_constraint(db, constraint, target)) .collect(); if constraint_tys.is_empty() { diff --git a/crates/ty_python_semantic/src/symbol.rs b/crates/ty_python_semantic/src/symbol.rs index 767d4392f29750..6a4fddf1311ba4 100644 --- a/crates/ty_python_semantic/src/symbol.rs +++ b/crates/ty_python_semantic/src/symbol.rs @@ -792,7 +792,7 @@ fn symbol_from_bindings_impl<'db>( } let binding_ty = binding_type(db, binding); - Some(narrowing_constraint.narrow(db, binding_ty, binding.symbol(db))) + Some(narrowing_constraint.narrow(db, binding_ty, binding.kind(db).target())) }, ); diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 0719800bfca919..36a6356b6aeea8 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -38,7 +38,7 @@ use ruff_db::diagnostic::{Annotation, DiagnosticId, Severity}; use ruff_db::files::File; use ruff_db::parsed::parsed_module; use ruff_python_ast::visitor::{walk_expr, Visitor}; -use ruff_python_ast::{self as ast, AnyNodeRef, ExprContext}; +use ruff_python_ast::{self as ast, AnyNodeRef, ExprContext, Identifier}; use ruff_text_size::{Ranged, TextRange}; use rustc_hash::{FxHashMap, FxHashSet}; use salsa; @@ -50,7 +50,7 @@ use crate::node_key::NodeKey; use crate::semantic_index::ast_ids::{HasScopedExpressionId, HasScopedUseId, ScopedExpressionId}; use crate::semantic_index::definition::{ AnnotatedAssignmentDefinitionKind, AssignmentDefinitionKind, ComprehensionDefinitionKind, - Definition, DefinitionKind, DefinitionNodeKey, ExceptHandlerDefinitionKind, + Definition, DefinitionKind, DefinitionNodeKey, DefinitionTarget, ExceptHandlerDefinitionKind, ForStmtDefinitionKind, TargetKind, WithItemDefinitionKind, }; use crate::semantic_index::expression::{Expression, ExpressionKind}; @@ -58,7 +58,7 @@ use crate::semantic_index::narrowing_constraints::ConstraintKey; use crate::semantic_index::symbol::{ FileScopeId, NodeWithScopeKind, NodeWithScopeRef, ScopeId, ScopeKind, }; -use crate::semantic_index::{semantic_index, EagerSnapshotResult, SemanticIndex}; +use crate::semantic_index::{semantic_index, use_def_map, EagerSnapshotResult, SemanticIndex}; use crate::symbol::{ builtins_module_scope, builtins_symbol, explicit_global_symbol, module_type_implicit_global_symbol, symbol, symbol_from_bindings, symbol_from_declarations, @@ -151,7 +151,7 @@ pub(crate) fn infer_definition_types<'db>( let file = definition.file(db); let _span = tracing::trace_span!( "infer_definition_types", - range = ?definition.kind(db).target_range(), + range = ?definition.kind(db).target().range(), ?file ) .entered(); @@ -190,7 +190,7 @@ pub(crate) fn infer_deferred_types<'db>( let _span = tracing::trace_span!( "infer_deferred_types", definition = ?definition.as_id(), - range = ?definition.kind(db).target_range(), + range = ?definition.kind(db).target().range(), ?file ) .entered(); @@ -5136,7 +5136,7 @@ impl<'db> TypeInferenceBuilder<'db> { /// Infer the type of a [`ast::ExprName`] expression, assuming a load context. fn infer_name_load(&mut self, name_node: &ast::ExprName) -> Type<'db> { let ast::ExprName { - range: _, + range, id: symbol_name, ctx: _, } = name_node; @@ -5153,10 +5153,10 @@ impl<'db> TypeInferenceBuilder<'db> { for (enclosing_scope_file_id, constraint_key) in constraint_keys { let use_def = self.index.use_def_map(*enclosing_scope_file_id); let constraints = use_def.narrowing_constraints_at_use(*constraint_key); - let symbol_table = self.index.symbol_table(*enclosing_scope_file_id); - let symbol = symbol_table.symbol_id_by_name(symbol_name).unwrap(); + let ident = Identifier::new(symbol_name.clone(), *range); + let target = DefinitionTarget::Ident(&ident); - ty = constraints.narrow(db, ty, symbol); + ty = constraints.narrow(db, ty, target); } ty }; @@ -5386,6 +5386,40 @@ impl<'db> TypeInferenceBuilder<'db> { } } + fn narrow<'r>(&self, target: impl Into>, ty: Type<'db>) -> Type<'db> { + fn root_name<'r>(expr: impl Into>) -> Option<&'r ast::ExprName> { + match expr.into() { + ast::ExprRef::Name(name) => Some(name), + ast::ExprRef::Attribute(attr) => root_name(&attr.value), + ast::ExprRef::Subscript(subscript) + if subscript.slice.is_string_literal_expr() + || subscript.slice.is_number_literal_expr() => + { + root_name(&subscript.value) + } + ast::ExprRef::Named(named) => root_name(&named.target), + _ => None, + } + } + + let target = target.into(); + + if let Some(name) = root_name(target) { + let db = self.db(); + let scope = self.scope(); + + let use_id = name.scoped_use_id(db, scope); + let use_def = use_def_map(db, scope); + let mut bindings = use_def.bindings_at_use(use_id); + let binding = bindings.next().unwrap(); + binding + .narrowing_constraint + .narrow(db, ty, DefinitionTarget::Expr(target)) + } else { + ty + } + } + /// Infer the type of a [`ast::ExprAttribute`] expression, assuming a load context. fn infer_attribute_load(&mut self, attribute: &ast::ExprAttribute) -> Type<'db> { let ast::ExprAttribute { @@ -5400,6 +5434,7 @@ impl<'db> TypeInferenceBuilder<'db> { value_type .member(db, &attr.id) + .map_type(|ty| self.narrow(attribute, ty)) .unwrap_with_diagnostic(|lookup_error| match lookup_error { LookupError::Unbound(_) => { let report_unresolved_attribute = self.is_reachable(attribute); @@ -6740,7 +6775,8 @@ impl<'db> TypeInferenceBuilder<'db> { } let slice_ty = self.infer_expression(slice); - self.infer_subscript_expression_types(value, value_ty, slice_ty) + let result_ty = self.infer_subscript_expression_types(value, value_ty, slice_ty); + self.narrow(subscript, result_ty) } fn infer_explicit_class_specialization( diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index b5a96c1a70dbd5..bff43b35a9f0f3 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -1,10 +1,11 @@ +use crate::ast_node_ref::AstNodeRef; use crate::semantic_index::ast_ids::HasScopedExpressionId; +use crate::semantic_index::definition::DefinitionTarget; use crate::semantic_index::expression::Expression; use crate::semantic_index::predicate::{ PatternPredicate, PatternPredicateKind, Predicate, PredicateNode, }; -use crate::semantic_index::symbol::{ScopeId, ScopedSymbolId, SymbolTable}; -use crate::semantic_index::symbol_table; +use crate::semantic_index::symbol::ScopeId; use crate::types::infer::infer_same_file_expression_type; use crate::types::{ infer_expression_types, IntersectionBuilder, KnownClass, SubclassOfType, Truthiness, Type, @@ -15,8 +16,8 @@ use itertools::Itertools; use ruff_python_ast as ast; use ruff_python_ast::{BoolOp, ExprBoolOp}; use rustc_hash::FxHashMap; +use salsa::Update; use std::collections::hash_map::Entry; -use std::sync::Arc; use super::UnionType; @@ -39,7 +40,7 @@ use super::UnionType; pub(crate) fn infer_narrowing_constraint<'db>( db: &'db dyn Db, predicate: Predicate<'db>, - symbol: ScopedSymbolId, + target: DefinitionTarget, ) -> Option> { let constraints = match predicate.node { PredicateNode::Expression(expression) => { @@ -59,7 +60,8 @@ pub(crate) fn infer_narrowing_constraint<'db>( PredicateNode::StarImportPlaceholder(_) => return None, }; if let Some(constraints) = constraints { - constraints.get(&symbol).copied() + let expr = NarrowingTarget::try_from_definition_target(target)?; + constraints.get(&expr).copied() } else { None } @@ -187,7 +189,89 @@ impl KnownConstraintFunction { } } -type NarrowingConstraints<'db> = FxHashMap>; +#[derive(Debug, Clone, PartialEq, Eq, Hash, Update)] +enum NarrowingTargetSegment { + /// A member access, e.g. `y` in `x.y`. + Member(ast::name::Name), + /// An integer index access, e.g. `1` in `x[1]`. + IntSubscript(ast::Int), + /// A string index access, e.g. `"foo"` in `x["foo"]`. + StringSubscript(String), +} + +/// An expression that is the target to type narrowing. +/// Names, member accesses, and subscriptions are supported (e.g., `x`, `x.y`, `x[0]`, `x.y[0].z`). +#[derive(Debug, Clone, PartialEq, Eq, Hash, Update)] +pub(crate) struct NarrowingTarget { + root_name: ast::name::Name, + segments: Vec, +} + +impl NarrowingTarget { + fn try_from_definition_target(node: DefinitionTarget) -> Option { + match node { + DefinitionTarget::Ident(id) => Some(NarrowingTarget { + root_name: id.id.clone(), + segments: vec![], + }), + DefinitionTarget::Expr(expr) => Self::try_from_expr(expr), + } + } + + fn try_from_node_ref(node_ref: &AstNodeRef) -> Option { + let expr: &ast::Expr = node_ref; + NarrowingTarget::try_from_expr(expr) + } + + fn try_from_expr<'r>(expr: impl Into>) -> Option { + match expr.into() { + ast::ExprRef::Name(ast::ExprName { id, .. }) => Some(NarrowingTarget { + root_name: id.clone(), + segments: vec![], + }), + ast::ExprRef::Attribute(ast::ExprAttribute { value, attr, .. }) => { + Self::try_from_expr(value).map(|mut expression| { + expression + .segments + .push(NarrowingTargetSegment::Member(attr.id().clone())); + expression + }) + } + ast::ExprRef::Subscript(ast::ExprSubscript { + value, + slice, + ctx: _, + range: _, + }) => match &**slice { + ast::Expr::NumberLiteral(number) if number.value.is_int() => { + Self::try_from_expr(value).map(|mut expression| { + expression + .segments + .push(NarrowingTargetSegment::IntSubscript( + number.value.as_int().unwrap().clone(), + )); + expression + }) + } + ast::Expr::StringLiteral(string) => { + Self::try_from_expr(value).map(|mut expression| { + expression + .segments + .push(NarrowingTargetSegment::StringSubscript( + string.value.to_str().to_string(), + )); + expression + }) + } + _ => None, + }, + ast::ExprRef::Named(named) => NarrowingTarget::try_from_expr(&named.target), + _ => None, + } + } +} + +type NarrowingConstraints<'db> = FxHashMap>; fn merge_constraints_and<'db>( into: &mut NarrowingConstraints<'db>, @@ -215,7 +299,7 @@ fn merge_constraints_or<'db>( db: &'db dyn Db, ) { for (key, value) in from { - match into.entry(*key) { + match into.entry(key.clone()) { Entry::Occupied(mut entry) => { *entry.get_mut() = UnionBuilder::new(db).add(*entry.get()).add(*value).build(); } @@ -237,17 +321,6 @@ fn negate_if<'db>(constraints: &mut NarrowingConstraints<'db>, db: &'db dyn Db, } } -fn expr_name(expr: &ast::Expr) -> Option<&ast::name::Name> { - match expr { - ast::Expr::Named(ast::ExprNamed { target, .. }) => match target.as_ref() { - ast::Expr::Name(ast::ExprName { id, .. }) => Some(id), - _ => None, - }, - ast::Expr::Name(ast::ExprName { id, .. }) => Some(id), - _ => None, - } -} - struct NarrowingConstraintsBuilder<'db> { db: &'db dyn Db, predicate: PredicateNode<'db>, @@ -297,7 +370,9 @@ impl<'db> NarrowingConstraintsBuilder<'db> { is_positive: bool, ) -> Option> { match expression_node { - ast::Expr::Name(name) => Some(self.evaluate_expr_name(name, is_positive)), + ast::Expr::Name(_) | ast::Expr::Attribute(_) | ast::Expr::Subscript(_) => { + self.evaluate_simple_expr(expression_node, is_positive) + } ast::Expr::Compare(expr_compare) => { self.evaluate_expr_compare(expr_compare, expression, is_positive) } @@ -344,10 +419,6 @@ impl<'db> NarrowingConstraintsBuilder<'db> { }) } - fn symbols(&self) -> Arc { - symbol_table(self.db, self.scope()) - } - fn scope(&self) -> ScopeId<'db> { match self.predicate { PredicateNode::Expression(expression) => expression.scope(self.db), @@ -356,21 +427,12 @@ impl<'db> NarrowingConstraintsBuilder<'db> { } } - #[track_caller] - fn expect_expr_name_symbol(&self, symbol: &str) -> ScopedSymbolId { - self.symbols() - .symbol_id_by_name(symbol) - .expect("We should always have a symbol for every `Name` node") - } - - fn evaluate_expr_name( + fn evaluate_simple_expr( &mut self, - expr_name: &ast::ExprName, + expr: &ast::Expr, is_positive: bool, - ) -> NarrowingConstraints<'db> { - let ast::ExprName { id, .. } = expr_name; - - let symbol = self.expect_expr_name_symbol(id); + ) -> Option> { + let target = NarrowingTarget::try_from_expr(expr)?; let ty = if is_positive { Type::AlwaysFalsy.negate(self.db) @@ -378,7 +440,7 @@ impl<'db> NarrowingConstraintsBuilder<'db> { Type::AlwaysTruthy.negate(self.db) }; - NarrowingConstraints::from_iter([(symbol, ty)]) + Some(NarrowingConstraints::from_iter([(target, ty)])) } fn evaluate_expr_named( @@ -386,11 +448,7 @@ impl<'db> NarrowingConstraintsBuilder<'db> { expr_named: &ast::ExprNamed, is_positive: bool, ) -> Option> { - if let ast::Expr::Name(expr_name) = expr_named.target.as_ref() { - Some(self.evaluate_expr_name(expr_name, is_positive)) - } else { - None - } + self.evaluate_simple_expr(&expr_named.target, is_positive) } fn evaluate_expr_eq(&mut self, lhs_ty: Type<'db>, rhs_ty: Type<'db>) -> Option> { @@ -582,7 +640,11 @@ impl<'db> NarrowingConstraintsBuilder<'db> { fn is_narrowing_target_candidate(expr: &ast::Expr) -> bool { matches!( expr, - ast::Expr::Name(_) | ast::Expr::Call(_) | ast::Expr::Named(_) + ast::Expr::Name(_) + | ast::Expr::Attribute(_) + | ast::Expr::Subscript(_) + | ast::Expr::Call(_) + | ast::Expr::Named(_) ) } @@ -627,13 +689,15 @@ impl<'db> NarrowingConstraintsBuilder<'db> { last_rhs_ty = Some(rhs_ty); match left { - ast::Expr::Name(_) | ast::Expr::Named(_) => { - if let Some(id) = expr_name(left) { - let symbol = self.expect_expr_name_symbol(id); + ast::Expr::Name(_) + | ast::Expr::Attribute(_) + | ast::Expr::Subscript(_) + | ast::Expr::Named(_) => { + if let Some(target) = NarrowingTarget::try_from_expr(left) { let op = if is_positive { *op } else { op.negate() }; if let Some(ty) = self.evaluate_expr_compare_op(lhs_ty, rhs_ty, op) { - constraints.insert(symbol, ty); + constraints.insert(target, ty); } } } @@ -655,9 +719,9 @@ impl<'db> NarrowingConstraintsBuilder<'db> { } }; - let id = match &**args { - [first] => match expr_name(first) { - Some(id) => id, + let target = match &**args { + [first] => match NarrowingTarget::try_from_expr(first) { + Some(target) => target, None => continue, }, _ => continue, @@ -680,9 +744,8 @@ impl<'db> NarrowingConstraintsBuilder<'db> { .into_class_literal() .is_some_and(|c| c.is_known(self.db, KnownClass::Type)) { - let symbol = self.expect_expr_name_symbol(id); constraints.insert( - symbol, + target, Type::instance(self.db, rhs_class.unknown_specialization(self.db)), ); } @@ -711,16 +774,14 @@ impl<'db> NarrowingConstraintsBuilder<'db> { Type::FunctionLiteral(function_type) if expr_call.arguments.keywords.is_empty() => { let function = function_type.known(self.db)?.into_constraint_function()?; - let (id, class_info) = match &*expr_call.arguments.args { - [first, class_info] => match expr_name(first) { - Some(id) => (id, class_info), + let (target, class_info) = match &*expr_call.arguments.args { + [first, class_info] => match NarrowingTarget::try_from_expr(first) { + Some(target) => (target, class_info), None => return None, }, _ => return None, }; - let symbol = self.expect_expr_name_symbol(id); - let class_info_ty = inference.expression_type(class_info.scoped_expression_id(self.db, scope)); @@ -728,7 +789,7 @@ impl<'db> NarrowingConstraintsBuilder<'db> { .generate_constraint(self.db, class_info_ty) .map(|constraint| { NarrowingConstraints::from_iter([( - symbol, + target, constraint.negate_if(self.db, !is_positive), )]) }) @@ -754,14 +815,14 @@ impl<'db> NarrowingConstraintsBuilder<'db> { subject: Expression<'db>, singleton: ast::Singleton, ) -> Option> { - let symbol = self.expect_expr_name_symbol(&subject.node_ref(self.db).as_name_expr()?.id); + let target = NarrowingTarget::try_from_node_ref(subject.node_ref(self.db))?; let ty = match singleton { ast::Singleton::None => Type::none(self.db), ast::Singleton::True => Type::BooleanLiteral(true), ast::Singleton::False => Type::BooleanLiteral(false), }; - Some(NarrowingConstraints::from_iter([(symbol, ty)])) + Some(NarrowingConstraints::from_iter([(target, ty)])) } fn evaluate_match_pattern_class( @@ -769,10 +830,11 @@ impl<'db> NarrowingConstraintsBuilder<'db> { subject: Expression<'db>, cls: Expression<'db>, ) -> Option> { - let symbol = self.expect_expr_name_symbol(&subject.node_ref(self.db).as_name_expr()?.id); + let target = NarrowingTarget::try_from_node_ref(subject.node_ref(self.db))?; + let ty = infer_same_file_expression_type(self.db, cls).to_instance(self.db)?; - Some(NarrowingConstraints::from_iter([(symbol, ty)])) + Some(NarrowingConstraints::from_iter([(target, ty)])) } fn evaluate_match_pattern_value( @@ -780,9 +842,10 @@ impl<'db> NarrowingConstraintsBuilder<'db> { subject: Expression<'db>, value: Expression<'db>, ) -> Option> { - let symbol = self.expect_expr_name_symbol(&subject.node_ref(self.db).as_name_expr()?.id); + let target = NarrowingTarget::try_from_node_ref(subject.node_ref(self.db))?; + let ty = infer_same_file_expression_type(self.db, value); - Some(NarrowingConstraints::from_iter([(symbol, ty)])) + Some(NarrowingConstraints::from_iter([(target, ty)])) } fn evaluate_match_pattern_or( From 9db97287e5fd90101da30c38c9bf4a0cdfbc34ce Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Sat, 26 Apr 2025 22:26:39 +0900 Subject: [PATCH 02/17] fix: use `try_scoped_use_id` --- .../src/semantic_index/ast_ids.rs | 20 +++++++++++++++++++ .../src/semantic_index/definition.rs | 2 +- crates/ty_python_semantic/src/types/infer.rs | 17 +++++++++------- 3 files changed, 31 insertions(+), 8 deletions(-) diff --git a/crates/ty_python_semantic/src/semantic_index/ast_ids.rs b/crates/ty_python_semantic/src/semantic_index/ast_ids.rs index bd5e93ada60df0..6ea89f1499cfc3 100644 --- a/crates/ty_python_semantic/src/semantic_index/ast_ids.rs +++ b/crates/ty_python_semantic/src/semantic_index/ast_ids.rs @@ -43,6 +43,10 @@ impl AstIds { fn use_id(&self, key: impl Into) -> ScopedUseId { self.uses_map[&key.into()] } + + fn try_use_id(&self, key: impl Into) -> Option { + self.uses_map.get(&key.into()).copied() + } } fn ast_ids<'db>(db: &'db dyn Db, scope: ScopeId) -> &'db AstIds { @@ -56,6 +60,7 @@ pub struct ScopedUseId; pub trait HasScopedUseId { /// Returns the ID that uniquely identifies the use in `scope`. fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId; + fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option; } impl HasScopedUseId for ast::Identifier { @@ -63,6 +68,11 @@ impl HasScopedUseId for ast::Identifier { let ast_ids = ast_ids(db, scope); ast_ids.use_id(self) } + + fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option { + let ast_ids = ast_ids(db, scope); + ast_ids.try_use_id(self) + } } impl HasScopedUseId for ast::ExprName { @@ -70,6 +80,11 @@ impl HasScopedUseId for ast::ExprName { let expression_ref = ExprRef::from(self); expression_ref.scoped_use_id(db, scope) } + + fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option { + let expression_ref = ExprRef::from(self); + expression_ref.try_scoped_use_id(db, scope) + } } impl HasScopedUseId for ast::ExprRef<'_> { @@ -77,6 +92,11 @@ impl HasScopedUseId for ast::ExprRef<'_> { let ast_ids = ast_ids(db, scope); ast_ids.use_id(*self) } + + fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option { + let ast_ids = ast_ids(db, scope); + ast_ids.try_use_id(*self) + } } /// Uniquely identifies an [`ast::Expr`] in a [`crate::semantic_index::symbol::FileScopeId`]. diff --git a/crates/ty_python_semantic/src/semantic_index/definition.rs b/crates/ty_python_semantic/src/semantic_index/definition.rs index c09afc3fae5020..fee44aa735cfa8 100644 --- a/crates/ty_python_semantic/src/semantic_index/definition.rs +++ b/crates/ty_python_semantic/src/semantic_index/definition.rs @@ -546,7 +546,7 @@ impl DefinitionTarget<'_> { /// /// A definition target would mainly be the node representing the symbol being defined i.e., /// [`ast::ExprName`] or [`ast::Identifier`] but could also be other nodes. - pub(crate) fn range(&self) -> TextRange { + pub(crate) fn range(self) -> TextRange { match self { Self::Ident(ident) => ident.range, Self::Expr(expr) => expr.range(), diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 36a6356b6aeea8..86b8ab5b036c11 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -5408,13 +5408,16 @@ impl<'db> TypeInferenceBuilder<'db> { let db = self.db(); let scope = self.scope(); - let use_id = name.scoped_use_id(db, scope); - let use_def = use_def_map(db, scope); - let mut bindings = use_def.bindings_at_use(use_id); - let binding = bindings.next().unwrap(); - binding - .narrowing_constraint - .narrow(db, ty, DefinitionTarget::Expr(target)) + if let Some(use_id) = name.try_scoped_use_id(db, scope) { + let use_def = use_def_map(db, scope); + let mut bindings = use_def.bindings_at_use(use_id); + let binding = bindings.next().unwrap(); + binding + .narrowing_constraint + .narrow(db, ty, DefinitionTarget::Expr(target)) + } else { + ty + } } else { ty } From 8a70f2a21c4dd70c9e646f73456961bfa8b874c0 Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Sat, 7 Jun 2025 13:42:34 +0900 Subject: [PATCH 03/17] add `PlaceExprWithFlags::{as_name, expect_name}` --- crates/ty_python_semantic/src/place.rs | 6 +++--- .../ty_python_semantic/src/semantic_index.rs | 2 +- .../src/semantic_index/place.rs | 18 ++++++++++++++---- .../semantic_index/visibility_constraints.rs | 1 - crates/ty_python_semantic/src/types/class.rs | 2 +- .../src/types/ide_support.rs | 8 ++++---- crates/ty_python_semantic/src/types/infer.rs | 3 +-- .../src/types/protocol_class.rs | 2 +- 8 files changed, 25 insertions(+), 17 deletions(-) diff --git a/crates/ty_python_semantic/src/place.rs b/crates/ty_python_semantic/src/place.rs index 95211d82221699..4e5ddf08b6e4b9 100644 --- a/crates/ty_python_semantic/src/place.rs +++ b/crates/ty_python_semantic/src/place.rs @@ -1011,7 +1011,7 @@ fn is_reexported(db: &dyn Db, definition: Definition<'_>) -> bool { return false; }; let table = place_table(db, definition.scope(db)); - let symbol_name = table.place_expr(definition.place(db)).expr.expect_name(); + let symbol_name = table.place_expr(definition.place(db)).expect_name(); all_names.contains(symbol_name) } @@ -1021,7 +1021,7 @@ mod implicit_globals { use crate::db::Db; use crate::place::PlaceAndQualifiers; use crate::semantic_index::place::PlaceExpr; - use crate::semantic_index::{place_table, use_def_map}; + use crate::semantic_index::{self, place_table, use_def_map}; use crate::types::{KnownClass, Type}; use super::{Place, PlaceFromDeclarationsResult, place_from_declarations}; @@ -1126,7 +1126,7 @@ mod implicit_globals { module_type_symbol_table .places() .filter(|place| place.is_declared() && place.is_name()) - .map(|symbol| symbol.expr.expect_name()) + .map(semantic_index::place::PlaceExprWithFlags::expect_name) .filter(|symbol_name| { !matches!(&***symbol_name, "__dict__" | "__getattr__" | "__init__") }) diff --git a/crates/ty_python_semantic/src/semantic_index.rs b/crates/ty_python_semantic/src/semantic_index.rs index 8fce599064a985..0b71130d1eb08c 100644 --- a/crates/ty_python_semantic/src/semantic_index.rs +++ b/crates/ty_python_semantic/src/semantic_index.rs @@ -605,7 +605,7 @@ mod tests { fn names(table: &PlaceTable) -> Vec { table .places() - .filter_map(|expr| Some(expr.expr.as_name()?.to_string())) + .filter_map(|expr| Some(expr.as_name()?.to_string())) .collect() } diff --git a/crates/ty_python_semantic/src/semantic_index/place.rs b/crates/ty_python_semantic/src/semantic_index/place.rs index 844199df449fc9..80d5d5872397d8 100644 --- a/crates/ty_python_semantic/src/semantic_index/place.rs +++ b/crates/ty_python_semantic/src/semantic_index/place.rs @@ -290,6 +290,14 @@ impl PlaceExprWithFlags { pub fn is_declared(&self) -> bool { self.flags.contains(PlaceFlags::IS_DECLARED) } + + pub(crate) fn as_name(&self) -> Option<&Name> { + self.expr.as_name() + } + + pub(crate) fn expect_name(&self) -> &Name { + self.expr.expect_name() + } } struct RootExprs<'e> { @@ -599,7 +607,7 @@ impl PlaceTable { .place_set .raw_entry() .from_hash(Self::hash_name(name), |id| { - self.place_expr(*id).expr.as_name().map(Name::as_str) == Some(name) + self.place_expr(*id).as_name().map(Name::as_str) == Some(name) })?; Some(*id) @@ -672,9 +680,11 @@ pub(super) struct PlaceTableBuilder { impl PlaceTableBuilder { pub(super) fn add_symbol(&mut self, name: Name) -> (ScopedPlaceId, bool) { let hash = PlaceTable::hash_name(&name); - let entry = self.table.place_set.raw_entry_mut().from_hash(hash, |id| { - self.table.places[*id].expr.as_name() == Some(&name) - }); + let entry = self + .table + .place_set + .raw_entry_mut() + .from_hash(hash, |id| self.table.places[*id].as_name() == Some(&name)); match entry { RawEntryMut::Occupied(entry) => (*entry.key(), false), diff --git a/crates/ty_python_semantic/src/semantic_index/visibility_constraints.rs b/crates/ty_python_semantic/src/semantic_index/visibility_constraints.rs index a221fa8c54ea4f..60db138c57da9a 100644 --- a/crates/ty_python_semantic/src/semantic_index/visibility_constraints.rs +++ b/crates/ty_python_semantic/src/semantic_index/visibility_constraints.rs @@ -657,7 +657,6 @@ impl VisibilityConstraints { let place_table = place_table(db, star_import.scope(db)); let symbol_name = place_table .place_expr(star_import.symbol_id(db)) - .expr .expect_name(); let referenced_file = star_import.referenced_file(db); diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index 6dc6f298b01133..344c2002379e27 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -1534,7 +1534,7 @@ impl<'db> ClassLiteral<'db> { let bindings = use_def.public_bindings(place_id); let default_ty = place_from_bindings(db, bindings).ignore_possibly_unbound(); - attributes.insert(place_expr.expr.expect_name().clone(), (attr_ty, default_ty)); + attributes.insert(place_expr.expect_name().clone(), (attr_ty, default_ty)); } } } diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index b8349bd6379f8c..e541e606c33266 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/crates/ty_python_semantic/src/types/ide_support.rs @@ -24,7 +24,7 @@ pub(crate) fn all_declarations_and_bindings<'db>( result .place .ignore_possibly_unbound() - .and_then(|_| table.place_expr(symbol_id).expr.as_name().cloned()) + .and_then(|_| table.place_expr(symbol_id).as_name().cloned()) }) }) .chain( @@ -33,7 +33,7 @@ pub(crate) fn all_declarations_and_bindings<'db>( .filter_map(move |(symbol_id, bindings)| { place_from_bindings(db, bindings) .ignore_possibly_unbound() - .and_then(|_| table.place_expr(symbol_id).expr.as_name().cloned()) + .and_then(|_| table.place_expr(symbol_id).as_name().cloned()) }), ) } @@ -134,7 +134,7 @@ impl AllMembers { let place_table = place_table(db, module_scope); for (symbol_id, _) in use_def_map.all_public_declarations() { - let Some(symbol_name) = place_table.place_expr(symbol_id).expr.as_name() else { + let Some(symbol_name) = place_table.place_expr(symbol_id).as_name() else { continue; }; if !imported_symbol(db, file, symbol_name, None) @@ -142,7 +142,7 @@ impl AllMembers { .is_unbound() { self.members - .insert(place_table.place_expr(symbol_id).expr.expect_name().clone()); + .insert(place_table.place_expr(symbol_id).expect_name().clone()); } } } diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 9186cb035359b6..62c8a366259131 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -1618,7 +1618,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let place_table = self.index.place_table(scope); let expr = place_table.place_expr(declaration.place(self.db())); if scope.is_global() && expr.is_name() { - module_type_implicit_global_symbol(self.db(), expr.expr.expect_name()) + module_type_implicit_global_symbol(self.db(), expr.expect_name()) } else { Place::Unbound.into() } @@ -4232,7 +4232,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let name = if let Some((star_import, symbol_table)) = star_import_info.as_ref() { symbol_table .place_expr(star_import.place_id()) - .expr .expect_name() } else { &alias.name.id diff --git a/crates/ty_python_semantic/src/types/protocol_class.rs b/crates/ty_python_semantic/src/types/protocol_class.rs index 2c1d54f5f694bf..b17864490e32e7 100644 --- a/crates/ty_python_semantic/src/types/protocol_class.rs +++ b/crates/ty_python_semantic/src/types/protocol_class.rs @@ -355,7 +355,7 @@ fn cached_protocol_interface<'db>( ) .filter_map(|(place_id, member, qualifiers)| { Some(( - place_table.place_expr(place_id).expr.as_name()?, + place_table.place_expr(place_id).as_name()?, member, qualifiers, )) From 06feb5cdc729b3d8baaf3781cb3606d4b71b0a2b Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Sat, 7 Jun 2025 15:46:39 +0900 Subject: [PATCH 04/17] use `constraint_keys` in attribute/subscript --- .../mdtest/narrow/conditionals/nested.md | 2 +- crates/ty_python_semantic/src/types/infer.rs | 43 +++++++++---------- 2 files changed, 22 insertions(+), 23 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md index 34891699f3453a..bd73933b6362f1 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md @@ -291,7 +291,7 @@ def f(x: str | Literal[1] | None): if a.x is not None: def _(): if a.x != 1: - reveal_type(a.x) # revealed: (Unknown & ~Literal[1]) | str + reveal_type(a.x) # revealed: (Unknown & ~Literal[1]) | str | None class D: if a.x != 1: diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 62c8a366259131..a0b67cf257f5c6 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -67,7 +67,7 @@ use crate::semantic_index::narrowing_constraints::ConstraintKey; use crate::semantic_index::place::{ FileScopeId, NodeWithScopeKind, NodeWithScopeRef, PlaceExpr, ScopeId, ScopeKind, ScopedPlaceId, }; -use crate::semantic_index::{EagerSnapshotResult, SemanticIndex, semantic_index, use_def_map}; +use crate::semantic_index::{EagerSnapshotResult, SemanticIndex, semantic_index}; use crate::types::call::{ Argument, Binding, Bindings, CallArgumentTypes, CallArguments, CallError, }; @@ -5776,11 +5776,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { mut ty: Type<'db>, constraint_keys: &[(FileScopeId, ConstraintKey)], ) -> Type<'db> { - tracing::warn!( - "5780: {expr}: {}, {}", - ty.display(self.db()), - constraint_keys.len() - ); let db = self.db(); for (enclosing_scope_file_id, constraint_key) in constraint_keys { let use_def = self.index.use_def_map(*enclosing_scope_file_id); @@ -5788,7 +5783,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { ty = constraints.narrow(db, ty, expr); } - tracing::warn!("-> {}", ty.display(self.db())); ty } @@ -6153,21 +6147,22 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - fn narrow<'r>(&self, target: impl Into>, ty: Type<'db>) -> Type<'db> { + fn narrow<'r>( + &self, + target: impl Into>, + ty: Type<'db>, + mut constraint_keys: Vec<(FileScopeId, ConstraintKey)>, + ) -> Type<'db> { let target = target.into(); if let Ok(place_expr) = PlaceExpr::try_from(target) { - let db = self.db(); - let scope = self.scope(); - - if let Some(use_id) = target.try_scoped_use_id(db, scope) { - let use_def = use_def_map(db, scope); - let mut bindings = use_def.bindings_at_use(use_id); - let binding = bindings.next().unwrap(); - binding.narrowing_constraint.narrow(db, ty, &place_expr) - } else { - ty + if let Some(use_id) = target.try_scoped_use_id(self.db(), self.scope()) { + constraint_keys.push(( + self.scope().file_scope_id(self.db()), + ConstraintKey::UseId(use_id), + )); } + self.narrow_with_applicable_constraints(&place_expr, ty, &constraint_keys) } else { ty } @@ -6184,6 +6179,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let value_type = self.infer_maybe_standalone_expression(value); let db = self.db(); + let mut constraint_keys = vec![]; // If `attribute` is a valid reference, we attempt type narrowing by assignment. if let Ok(place_expr) = PlaceExpr::try_from(attribute) { @@ -6195,8 +6191,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .ignore_possibly_unbound() .is_none_or(|ty| !ty.may_be_data_descriptor(db)) { - let (resolved, _) = + let (resolved, keys) = self.infer_place_load(&place_expr, ast::ExprRef::Attribute(attribute)); + constraint_keys.extend(keys); if let Place::Type(ty, Boundness::Bound) = resolved.place { return ty; } @@ -6205,7 +6202,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { value_type .member(db, &attr.id) - .map_type(|ty| self.narrow(attribute, ty)) + .map_type(|ty| self.narrow(attribute, ty, constraint_keys)) .unwrap_with_diagnostic(|lookup_error| match lookup_error { LookupError::Unbound(_) => { let report_unresolved_attribute = self.is_reachable(attribute); @@ -7694,6 +7691,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } = subscript; let db = self.db(); let value_ty = self.infer_expression(value); + let mut constraint_keys = vec![]; // If `value` is a valid reference, we attempt type narrowing by assignment. if !value_ty.is_unknown() { @@ -7726,8 +7724,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .zip(safe_mutable_class.generic_origin(db)) .is_some_and(|(l, r)| l == r) }) { - let (place, _) = + let (place, keys) = self.infer_place_load(&expr, ast::ExprRef::Subscript(subscript)); + constraint_keys.extend(keys); if let Place::Type(ty, Boundness::Bound) = place.place { self.infer_expression(slice); return ty; @@ -7762,7 +7761,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let slice_ty = self.infer_expression(slice); let result_ty = self.infer_subscript_expression_types(value, value_ty, slice_ty); - self.narrow(subscript, result_ty) + self.narrow(subscript, result_ty, constraint_keys) } fn infer_explicit_class_specialization( From 0c8996c501650215a55fc73348a689e10fe3fbff Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Sat, 7 Jun 2025 18:12:47 +0900 Subject: [PATCH 05/17] use `ScopedPlaceId` as `NarrowingConstraints` key --- crates/ty_python_semantic/src/place.rs | 4 +- .../src/semantic_index/place.rs | 6 +-- .../src/semantic_index/use_def.rs | 8 ++-- crates/ty_python_semantic/src/types/infer.rs | 8 ++-- crates/ty_python_semantic/src/types/narrow.rs | 45 +++++++++++++------ 5 files changed, 44 insertions(+), 27 deletions(-) diff --git a/crates/ty_python_semantic/src/place.rs b/crates/ty_python_semantic/src/place.rs index 4e5ddf08b6e4b9..a3b4fb038a78e6 100644 --- a/crates/ty_python_semantic/src/place.rs +++ b/crates/ty_python_semantic/src/place.rs @@ -863,9 +863,7 @@ fn place_from_bindings_impl<'db>( } let binding_ty = binding_type(db, binding); - let place_table = place_table(db, binding.scope(db)); - let place_expr = place_table.place_expr(binding.place(db)); - Some(narrowing_constraint.narrow(db, binding_ty, &place_expr.expr)) + Some(narrowing_constraint.narrow(db, binding_ty, binding.place(db))) }, ); diff --git a/crates/ty_python_semantic/src/semantic_index/place.rs b/crates/ty_python_semantic/src/semantic_index/place.rs index 80d5d5872397d8..cbbdb3e8a8bb71 100644 --- a/crates/ty_python_semantic/src/semantic_index/place.rs +++ b/crates/ty_python_semantic/src/semantic_index/place.rs @@ -18,7 +18,7 @@ use crate::node_key::NodeKey; use crate::semantic_index::visibility_constraints::ScopedVisibilityConstraintId; use crate::semantic_index::{PlaceSet, SemanticIndex, semantic_index}; -#[derive(Debug, Clone, PartialEq, Eq, Hash, salsa::Update)] +#[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum PlaceExprSubSegment { /// A member access, e.g. `.y` in `x.y` Member(ast::name::Name), @@ -38,7 +38,7 @@ impl PlaceExprSubSegment { } /// An expression that can be the target of a `Definition`. -#[derive(Clone, Eq, PartialEq, Debug, Hash, salsa::Update)] +#[derive(Eq, PartialEq, Debug)] pub struct PlaceExpr { root_name: Name, sub_segments: SmallVec<[PlaceExprSubSegment; 1]>, @@ -325,7 +325,7 @@ bitflags! { /// /// See the doc-comment at the top of [`super::use_def`] for explanations of what it /// means for a place to be *bound* as opposed to *declared*. - #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)] + #[derive(Copy, Clone, Debug, Eq, PartialEq)] struct PlaceFlags: u8 { const IS_USED = 1 << 0; const IS_BOUND = 1 << 1; diff --git a/crates/ty_python_semantic/src/semantic_index/use_def.rs b/crates/ty_python_semantic/src/semantic_index/use_def.rs index 4811ef64d1c7cb..eab70cfb57e9ae 100644 --- a/crates/ty_python_semantic/src/semantic_index/use_def.rs +++ b/crates/ty_python_semantic/src/semantic_index/use_def.rs @@ -278,9 +278,7 @@ use crate::semantic_index::definition::{Definition, DefinitionState}; use crate::semantic_index::narrowing_constraints::{ ConstraintKey, NarrowingConstraints, NarrowingConstraintsBuilder, NarrowingConstraintsIterator, }; -use crate::semantic_index::place::{ - FileScopeId, PlaceExpr, PlaceExprWithFlags, ScopeKind, ScopedPlaceId, -}; +use crate::semantic_index::place::{FileScopeId, PlaceExprWithFlags, ScopeKind, ScopedPlaceId}; use crate::semantic_index::predicate::{ Predicate, Predicates, PredicatesBuilder, ScopedPredicateId, StarImportPlaceholderPredicate, }; @@ -595,10 +593,10 @@ impl<'db> ConstraintsIterator<'_, 'db> { self, db: &'db dyn crate::Db, base_ty: Type<'db>, - place_expr: &PlaceExpr, + place: ScopedPlaceId, ) -> Type<'db> { let constraint_tys: Vec<_> = self - .filter_map(|constraint| infer_narrowing_constraint(db, constraint, place_expr)) + .filter_map(|constraint| infer_narrowing_constraint(db, constraint, place)) .collect(); if constraint_tys.is_empty() { diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index a0b67cf257f5c6..9fa197c15e7a94 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -5780,8 +5780,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { for (enclosing_scope_file_id, constraint_key) in constraint_keys { let use_def = self.index.use_def_map(*enclosing_scope_file_id); let constraints = use_def.narrowing_constraints_at_use(*constraint_key); + let place_table = self.index.place_table(*enclosing_scope_file_id); + let place = place_table.place_id_by_expr(expr).unwrap(); - ty = constraints.narrow(db, ty, expr); + ty = constraints.narrow(db, ty, place); } ty } @@ -6193,10 +6195,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { { let (resolved, keys) = self.infer_place_load(&place_expr, ast::ExprRef::Attribute(attribute)); - constraint_keys.extend(keys); if let Place::Type(ty, Boundness::Bound) = resolved.place { return ty; } + constraint_keys.extend(keys); } } @@ -7726,11 +7728,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { }) { let (place, keys) = self.infer_place_load(&expr, ast::ExprRef::Subscript(subscript)); - constraint_keys.extend(keys); if let Place::Type(ty, Boundness::Bound) = place.place { self.infer_expression(slice); return ty; } + constraint_keys.extend(keys); } } } diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index 1aae650c1d7402..502f978fb3cf58 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -1,7 +1,8 @@ use crate::Db; use crate::semantic_index::ast_ids::HasScopedExpressionId; use crate::semantic_index::expression::Expression; -use crate::semantic_index::place::{PlaceExpr, ScopeId}; +use crate::semantic_index::place::{PlaceExpr, PlaceTable, ScopeId, ScopedPlaceId}; +use crate::semantic_index::place_table; use crate::semantic_index::predicate::{ PatternPredicate, PatternPredicateKind, Predicate, PredicateNode, }; @@ -42,7 +43,7 @@ use super::UnionType; pub(crate) fn infer_narrowing_constraint<'db>( db: &'db dyn Db, predicate: Predicate<'db>, - place: &PlaceExpr, + place: ScopedPlaceId, ) -> Option> { let constraints = match predicate.node { PredicateNode::Expression(expression) => { @@ -62,7 +63,7 @@ pub(crate) fn infer_narrowing_constraint<'db>( PredicateNode::StarImportPlaceholder(_) => return None, }; if let Some(constraints) = constraints { - constraints.get(place).copied() + constraints.get(&place).copied() } else { None } @@ -196,7 +197,7 @@ impl ClassInfoConstraintFunction { } } -type NarrowingConstraints<'db> = FxHashMap>; +type NarrowingConstraints<'db> = FxHashMap>; fn merge_constraints_and<'db>( into: &mut NarrowingConstraints<'db>, @@ -224,7 +225,7 @@ fn merge_constraints_or<'db>( db: &'db dyn Db, ) { for (key, value) in from { - match into.entry(key.clone()) { + match into.entry(*key) { Entry::Occupied(mut entry) => { *entry.get_mut() = UnionBuilder::new(db).add(*entry.get()).add(*value).build(); } @@ -351,6 +352,10 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { }) } + fn places(&self) -> &'db PlaceTable { + place_table(self.db, self.scope()) + } + fn scope(&self) -> ScopeId<'db> { match self.predicate { PredicateNode::Expression(expression) => expression.scope(self.db), @@ -359,12 +364,20 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { } } + #[track_caller] + fn expect_place(&self, place_expr: &PlaceExpr) -> ScopedPlaceId { + self.places() + .place_id_by_expr(place_expr) + .expect("We should always have a place for every `PlaceExpr`") + } + fn evaluate_simple_expr( &mut self, expr: &ast::Expr, is_positive: bool, ) -> Option> { let target = PlaceExpr::try_from(expr).ok()?; + let place = self.expect_place(&target); let ty = if is_positive { Type::AlwaysFalsy.negate(self.db) @@ -372,7 +385,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { Type::AlwaysTruthy.negate(self.db) }; - Some(NarrowingConstraints::from_iter([(target, ty)])) + Some(NarrowingConstraints::from_iter([(place, ty)])) } fn evaluate_expr_named( @@ -629,7 +642,8 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { let op = if is_positive { *op } else { op.negate() }; if let Some(ty) = self.evaluate_expr_compare_op(lhs_ty, rhs_ty, op) { - constraints.insert(target, ty); + let place = self.expect_place(&target); + constraints.insert(place, ty); } } } @@ -676,8 +690,9 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { .into_class_literal() .is_some_and(|c| c.is_known(self.db, KnownClass::Type)) { + let place = self.expect_place(&target); constraints.insert( - target, + place, Type::instance(self.db, rhs_class.unknown_specialization(self.db)), ); } @@ -713,6 +728,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { return None; }; let function = function_type.known(self.db)?; + let place = self.expect_place(&target); if function == KnownFunction::HasAttr { let attr = inference @@ -730,7 +746,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { ); return Some(NarrowingConstraints::from_iter([( - target, + place, constraint.negate_if(self.db, !is_positive), )])); } @@ -744,7 +760,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { .generate_constraint(self.db, class_info_ty) .map(|constraint| { NarrowingConstraints::from_iter([( - target, + place, constraint.negate_if(self.db, !is_positive), )]) }) @@ -771,13 +787,14 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { singleton: ast::Singleton, ) -> Option> { let target = PlaceExpr::try_from(subject.node_ref(self.db, self.module)).ok()?; + let place = self.expect_place(&target); let ty = match singleton { ast::Singleton::None => Type::none(self.db), ast::Singleton::True => Type::BooleanLiteral(true), ast::Singleton::False => Type::BooleanLiteral(false), }; - Some(NarrowingConstraints::from_iter([(target, ty)])) + Some(NarrowingConstraints::from_iter([(place, ty)])) } fn evaluate_match_pattern_class( @@ -786,10 +803,11 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { cls: Expression<'db>, ) -> Option> { let target = PlaceExpr::try_from(subject.node_ref(self.db, self.module)).ok()?; + let place = self.expect_place(&target); let ty = infer_same_file_expression_type(self.db, cls, self.module).to_instance(self.db)?; - Some(NarrowingConstraints::from_iter([(target, ty)])) + Some(NarrowingConstraints::from_iter([(place, ty)])) } fn evaluate_match_pattern_value( @@ -798,9 +816,10 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { value: Expression<'db>, ) -> Option> { let target = PlaceExpr::try_from(subject.node_ref(self.db, self.module)).ok()?; + let place = self.expect_place(&target); let ty = infer_same_file_expression_type(self.db, value, self.module); - Some(NarrowingConstraints::from_iter([(target, ty)])) + Some(NarrowingConstraints::from_iter([(place, ty)])) } fn evaluate_match_pattern_or( From f29893a7d5a18c088ea35691b1f288a1d1ff980b Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Sat, 7 Jun 2025 21:19:21 +0900 Subject: [PATCH 06/17] Don't convert `ast::ExprNamed` with `PlaceExpr::try_from(ast::Expr)` --- .../src/semantic_index/place.rs | 2 - crates/ty_python_semantic/src/types/narrow.rs | 42 +++++++++++-------- 2 files changed, 24 insertions(+), 20 deletions(-) diff --git a/crates/ty_python_semantic/src/semantic_index/place.rs b/crates/ty_python_semantic/src/semantic_index/place.rs index cbbdb3e8a8bb71..69f1da96ffb812 100644 --- a/crates/ty_python_semantic/src/semantic_index/place.rs +++ b/crates/ty_python_semantic/src/semantic_index/place.rs @@ -143,7 +143,6 @@ impl TryFrom<&ast::Expr> for PlaceExpr { ast::Expr::Name(name) => Ok(PlaceExpr::name(name.id.clone())), ast::Expr::Attribute(attr) => PlaceExpr::try_from(attr), ast::Expr::Subscript(subscript) => PlaceExpr::try_from(subscript), - ast::Expr::Named(named) => PlaceExpr::try_from(named.target.as_ref()), _ => Err(()), } } @@ -157,7 +156,6 @@ impl TryFrom> for PlaceExpr { ast::ExprRef::Name(name) => Ok(PlaceExpr::name(name.id.clone())), ast::ExprRef::Attribute(attr) => PlaceExpr::try_from(attr), ast::ExprRef::Subscript(subscript) => PlaceExpr::try_from(subscript), - ast::ExprRef::Named(named) => PlaceExpr::try_from(named.target.as_ref()), _ => Err(()), } } diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index 502f978fb3cf58..be1e3d2fe5a59b 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -247,6 +247,16 @@ fn negate_if<'db>(constraints: &mut NarrowingConstraints<'db>, db: &'db dyn Db, } } +fn place_expr(expr: &ast::Expr) -> Option { + match expr { + ast::Expr::Name(name) => Some(PlaceExpr::name(name.id.clone())), + ast::Expr::Attribute(attr) => PlaceExpr::try_from(attr).ok(), + ast::Expr::Subscript(subscript) => PlaceExpr::try_from(subscript).ok(), + ast::Expr::Named(named) => PlaceExpr::try_from(named.target.as_ref()).ok(), + _ => None, + } +} + struct NarrowingConstraintsBuilder<'db, 'ast> { db: &'db dyn Db, module: &'ast ParsedModuleRef, @@ -376,7 +386,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { expr: &ast::Expr, is_positive: bool, ) -> Option> { - let target = PlaceExpr::try_from(expr).ok()?; + let target = place_expr(expr)?; let place = self.expect_place(&target); let ty = if is_positive { @@ -638,11 +648,11 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { | ast::Expr::Attribute(_) | ast::Expr::Subscript(_) | ast::Expr::Named(_) => { - if let Ok(target) = PlaceExpr::try_from(left) { + if let Some(left) = place_expr(left) { let op = if is_positive { *op } else { op.negate() }; if let Some(ty) = self.evaluate_expr_compare_op(lhs_ty, rhs_ty, op) { - let place = self.expect_place(&target); + let place = self.expect_place(&left); constraints.insert(place, ty); } } @@ -666,9 +676,9 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { }; let target = match &**args { - [first] => match PlaceExpr::try_from(first) { - Ok(target) => target, - Err(()) => continue, + [first] => match place_expr(first) { + Some(target) => target, + None => continue, }, _ => continue, }; @@ -719,16 +729,12 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { // and `issubclass`, for example `isinstance(x, str | (int | float))`. match callable_ty { Type::FunctionLiteral(function_type) if expr_call.arguments.keywords.is_empty() => { - // let function = function_type.known(self.db)?.into_constraint_function()?; - let [first_arg, second_arg] = &*expr_call.arguments.args else { return None; }; - let Ok(target) = PlaceExpr::try_from(first_arg) else { - return None; - }; + let first_arg = place_expr(first_arg)?; let function = function_type.known(self.db)?; - let place = self.expect_place(&target); + let place = self.expect_place(&first_arg); if function == KnownFunction::HasAttr { let attr = inference @@ -786,8 +792,8 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { subject: Expression<'db>, singleton: ast::Singleton, ) -> Option> { - let target = PlaceExpr::try_from(subject.node_ref(self.db, self.module)).ok()?; - let place = self.expect_place(&target); + let subject = place_expr(subject.node_ref(self.db, self.module))?; + let place = self.expect_place(&subject); let ty = match singleton { ast::Singleton::None => Type::none(self.db), @@ -802,8 +808,8 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { subject: Expression<'db>, cls: Expression<'db>, ) -> Option> { - let target = PlaceExpr::try_from(subject.node_ref(self.db, self.module)).ok()?; - let place = self.expect_place(&target); + let subject = place_expr(subject.node_ref(self.db, self.module))?; + let place = self.expect_place(&subject); let ty = infer_same_file_expression_type(self.db, cls, self.module).to_instance(self.db)?; @@ -815,8 +821,8 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> { subject: Expression<'db>, value: Expression<'db>, ) -> Option> { - let target = PlaceExpr::try_from(subject.node_ref(self.db, self.module)).ok()?; - let place = self.expect_place(&target); + let subject = place_expr(subject.node_ref(self.db, self.module))?; + let place = self.expect_place(&subject); let ty = infer_same_file_expression_type(self.db, value, self.module); Some(NarrowingConstraints::from_iter([(place, ty)])) From f74673095b4edc4741d67299c46edfd850b533d7 Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Sat, 7 Jun 2025 21:51:52 +0900 Subject: [PATCH 07/17] add `PlaceExprRef` --- .../src/semantic_index/place.rs | 68 ++++++++++++++----- 1 file changed, 51 insertions(+), 17 deletions(-) diff --git a/crates/ty_python_semantic/src/semantic_index/place.rs b/crates/ty_python_semantic/src/semantic_index/place.rs index 69f1da96ffb812..00160b0530129c 100644 --- a/crates/ty_python_semantic/src/semantic_index/place.rs +++ b/crates/ty_python_semantic/src/semantic_index/place.rs @@ -208,11 +208,9 @@ impl PlaceExpr { .is_some_and(|last| last.as_member().is_some()) } - // TODO: Ideally this would iterate PlaceSegments instead of RootExprs, both to reduce - // allocation and to avoid having both flagged and non-flagged versions of PlaceExprs. fn root_exprs(&self) -> RootExprs<'_> { RootExprs { - expr: self, + expr_ref: self.into(), len: self.sub_segments.len(), } } @@ -298,22 +296,49 @@ impl PlaceExprWithFlags { } } +#[derive(Clone, Copy, Eq, PartialEq, Debug)] +pub struct PlaceExprRef<'a> { + pub(crate) root_name: &'a Name, + pub(crate) sub_segments: &'a [PlaceExprSubSegment], +} + +impl PartialEq for PlaceExprRef<'_> { + fn eq(&self, other: &PlaceExpr) -> bool { + self.root_name == &other.root_name && self.sub_segments == &other.sub_segments[..] + } +} + +impl PartialEq> for PlaceExpr { + fn eq(&self, other: &PlaceExprRef<'_>) -> bool { + &self.root_name == other.root_name && &self.sub_segments[..] == other.sub_segments + } +} + +impl<'e> From<&'e PlaceExpr> for PlaceExprRef<'e> { + fn from(expr: &'e PlaceExpr) -> Self { + PlaceExprRef { + root_name: &expr.root_name, + sub_segments: &expr.sub_segments, + } + } +} + struct RootExprs<'e> { - expr: &'e PlaceExpr, + expr_ref: PlaceExprRef<'e>, len: usize, } -impl Iterator for RootExprs<'_> { - type Item = PlaceExpr; +impl<'e> Iterator for RootExprs<'e> { + type Item = PlaceExprRef<'e>; fn next(&mut self) -> Option { if self.len == 0 { return None; } self.len -= 1; - Some(PlaceExpr { - root_name: self.expr.root_name.clone(), - sub_segments: self.expr.sub_segments[..self.len].iter().cloned().collect(), + Some(PlaceExprRef { + root_name: self.expr_ref.root_name, + sub_segments: &self.expr_ref.sub_segments[..self.len], }) } } @@ -565,7 +590,7 @@ impl PlaceTable { ) -> impl Iterator { place_expr .root_exprs() - .filter_map(|place_expr| self.place_by_expr(&place_expr)) + .filter_map(|place_expr| self.place_by_expr(place_expr)) } #[expect(unused)] @@ -594,7 +619,10 @@ impl PlaceTable { } /// Returns the flagged place. - pub(crate) fn place_by_expr(&self, place_expr: &PlaceExpr) -> Option<&PlaceExprWithFlags> { + pub(crate) fn place_by_expr<'e>( + &self, + place_expr: impl Into>, + ) -> Option<&PlaceExprWithFlags> { let id = self.place_id_by_expr(place_expr)?; Some(self.place_expr(id)) } @@ -612,12 +640,16 @@ impl PlaceTable { } /// Returns the [`ScopedPlaceId`] of the place expression. - pub(crate) fn place_id_by_expr(&self, place_expr: &PlaceExpr) -> Option { + pub(crate) fn place_id_by_expr<'e>( + &self, + place_expr: impl Into>, + ) -> Option { + let place_expr = place_expr.into(); let (id, ()) = self .place_set .raw_entry() .from_hash(Self::hash_place_expr(place_expr), |id| { - &self.place_expr(*id).expr == place_expr + self.place_expr(*id).expr == place_expr })?; Some(*id) @@ -635,10 +667,12 @@ impl PlaceTable { hasher.finish() } - fn hash_place_expr(place_expr: &PlaceExpr) -> u64 { + fn hash_place_expr<'e>(place_expr: impl Into>) -> u64 { + let place_expr = place_expr.into(); + let mut hasher = FxHasher::default(); - place_expr.root_name().as_str().hash(&mut hasher); - for segment in &place_expr.sub_segments { + place_expr.root_name.as_str().hash(&mut hasher); + for segment in place_expr.sub_segments { match segment { PlaceExprSubSegment::Member(name) => name.hash(&mut hasher), PlaceExprSubSegment::IntSubscript(int) => int.hash(&mut hasher), @@ -718,7 +752,7 @@ impl PlaceTableBuilder { let new_id = self.associated_place_ids.push(vec![]); debug_assert_eq!(new_id, id); for root in self.table.places[id].expr.root_exprs() { - if let Some(root_id) = self.table.place_id_by_expr(&root) { + if let Some(root_id) = self.table.place_id_by_expr(root) { self.associated_place_ids[root_id].push(id); } } From 0cabb63c4fd2d55c0b64a65c4cef33629c79fdb7 Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Mon, 9 Jun 2025 15:39:39 +0900 Subject: [PATCH 08/17] remove `try_scoped_use_id` --- .../src/semantic_index/ast_ids.rs | 30 ---------- crates/ty_python_semantic/src/types/infer.rs | 56 +++++++++---------- 2 files changed, 25 insertions(+), 61 deletions(-) diff --git a/crates/ty_python_semantic/src/semantic_index/ast_ids.rs b/crates/ty_python_semantic/src/semantic_index/ast_ids.rs index 9ccb9ffefdbbb4..191be73c23244d 100644 --- a/crates/ty_python_semantic/src/semantic_index/ast_ids.rs +++ b/crates/ty_python_semantic/src/semantic_index/ast_ids.rs @@ -43,10 +43,6 @@ impl AstIds { fn use_id(&self, key: impl Into) -> ScopedUseId { self.uses_map[&key.into()] } - - fn try_use_id(&self, key: impl Into) -> Option { - self.uses_map.get(&key.into()).copied() - } } fn ast_ids<'db>(db: &'db dyn Db, scope: ScopeId) -> &'db AstIds { @@ -60,7 +56,6 @@ pub struct ScopedUseId; pub trait HasScopedUseId { /// Returns the ID that uniquely identifies the use in `scope`. fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId; - fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option; } impl HasScopedUseId for ast::Identifier { @@ -68,11 +63,6 @@ impl HasScopedUseId for ast::Identifier { let ast_ids = ast_ids(db, scope); ast_ids.use_id(self) } - - fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option { - let ast_ids = ast_ids(db, scope); - ast_ids.try_use_id(self) - } } impl HasScopedUseId for ast::ExprName { @@ -80,11 +70,6 @@ impl HasScopedUseId for ast::ExprName { let expression_ref = ExprRef::from(self); expression_ref.scoped_use_id(db, scope) } - - fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option { - let expression_ref = ExprRef::from(self); - expression_ref.try_scoped_use_id(db, scope) - } } impl HasScopedUseId for ast::ExprAttribute { @@ -92,11 +77,6 @@ impl HasScopedUseId for ast::ExprAttribute { let expression_ref = ExprRef::from(self); expression_ref.scoped_use_id(db, scope) } - - fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option { - let expression_ref = ExprRef::from(self); - expression_ref.try_scoped_use_id(db, scope) - } } impl HasScopedUseId for ast::ExprSubscript { @@ -104,11 +84,6 @@ impl HasScopedUseId for ast::ExprSubscript { let expression_ref = ExprRef::from(self); expression_ref.scoped_use_id(db, scope) } - - fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option { - let expression_ref = ExprRef::from(self); - expression_ref.try_scoped_use_id(db, scope) - } } impl HasScopedUseId for ast::ExprRef<'_> { @@ -116,11 +91,6 @@ impl HasScopedUseId for ast::ExprRef<'_> { let ast_ids = ast_ids(db, scope); ast_ids.use_id(*self) } - - fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option { - let ast_ids = ast_ids(db, scope); - ast_ids.try_use_id(*self) - } } /// Uniquely identifies an [`ast::Expr`] in a [`crate::semantic_index::place::FileScopeId`]. diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 9fa197c15e7a94..d25e3e66aa8cb7 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -5770,7 +5770,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } // Perform narrowing with applicable constraints between the current scope and the enclosing scope. - fn narrow_with_applicable_constraints( + fn narrow_place_with_applicable_constraints( &self, expr: &PlaceExpr, mut ty: Type<'db>, @@ -5805,7 +5805,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // These are looked up as attributes on `types.ModuleType`. .or_fall_back_to(db, || { module_type_implicit_global_symbol(db, symbol_name).map_type(|ty| { - self.narrow_with_applicable_constraints(&expr, ty, &constraint_keys) + self.narrow_place_with_applicable_constraints(&expr, ty, &constraint_keys) }) }) // Not found in globals? Fallback to builtins @@ -5877,7 +5877,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - /// Infer the type of a place expression, assuming a load context. + /// Infer the type of a place expression from definitions, assuming a load context. + /// This method also returns the [`ConstraintKey`]s for each scope associated with `expr`, + /// which is used to narrow by condition rather than by assignment. fn infer_place_load( &self, expr: &PlaceExpr, @@ -5890,6 +5892,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let mut constraint_keys = vec![]; let (local_scope_place, use_id) = self.infer_local_place_load(expr, expr_ref); + if let Some(use_id) = use_id { + constraint_keys.push((file_scope_id, ConstraintKey::UseId(use_id))); + } let place = PlaceAndQualifiers::from(local_scope_place).or_fall_back_to(db, || { let has_bindings_in_this_scope = match place_table.place_by_expr(expr) { @@ -5948,10 +5953,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - if let Some(use_id) = use_id { - constraint_keys.push((file_scope_id, ConstraintKey::UseId(use_id))); - } - // Walk up parent scopes looking for a possible enclosing scope that may have a // definition of this name visible to us (would be `LOAD_DEREF` at runtime.) // Note that we skip the scope containing the use that we are resolving, since we @@ -5996,7 +5997,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } return place_from_bindings(db, bindings) .map_type(|ty| { - self.narrow_with_applicable_constraints( + self.narrow_place_with_applicable_constraints( expr, ty, &constraint_keys, @@ -6041,7 +6042,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // isn't bound in that scope, we should get an unbound name, not continue // falling back to other scopes / globals / builtins. return place(db, enclosing_scope_id, expr).map_type(|ty| { - self.narrow_with_applicable_constraints(expr, ty, &constraint_keys) + self.narrow_place_with_applicable_constraints(expr, ty, &constraint_keys) }); } } @@ -6068,7 +6069,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { EagerSnapshotResult::FoundBindings(bindings) => { return place_from_bindings(db, bindings) .map_type(|ty| { - self.narrow_with_applicable_constraints( + self.narrow_place_with_applicable_constraints( expr, ty, &constraint_keys, @@ -6089,7 +6090,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { }; explicit_global_symbol(db, self.file(), name).map_type(|ty| { - self.narrow_with_applicable_constraints(expr, ty, &constraint_keys) + self.narrow_place_with_applicable_constraints(expr, ty, &constraint_keys) }) }) }); @@ -6149,24 +6150,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - fn narrow<'r>( + fn narrow_expr_with_applicable_constraints<'r>( &self, target: impl Into>, - ty: Type<'db>, - mut constraint_keys: Vec<(FileScopeId, ConstraintKey)>, + target_ty: Type<'db>, + constraint_keys: &[(FileScopeId, ConstraintKey)], ) -> Type<'db> { let target = target.into(); if let Ok(place_expr) = PlaceExpr::try_from(target) { - if let Some(use_id) = target.try_scoped_use_id(self.db(), self.scope()) { - constraint_keys.push(( - self.scope().file_scope_id(self.db()), - ConstraintKey::UseId(use_id), - )); - } - self.narrow_with_applicable_constraints(&place_expr, ty, &constraint_keys) + self.narrow_place_with_applicable_constraints(&place_expr, target_ty, constraint_keys) } else { - ty + target_ty } } @@ -6183,8 +6178,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let db = self.db(); let mut constraint_keys = vec![]; - // If `attribute` is a valid reference, we attempt type narrowing by assignment. if let Ok(place_expr) = PlaceExpr::try_from(attribute) { + let (resolved, keys) = + self.infer_place_load(&place_expr, ast::ExprRef::Attribute(attribute)); + constraint_keys.extend(keys); let member = value_type.class_member(db, attr.id.clone()); // If the member is a data descriptor, the value most recently assigned // to the attribute may not necessarily be obtained here. @@ -6193,18 +6190,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .ignore_possibly_unbound() .is_none_or(|ty| !ty.may_be_data_descriptor(db)) { - let (resolved, keys) = - self.infer_place_load(&place_expr, ast::ExprRef::Attribute(attribute)); if let Place::Type(ty, Boundness::Bound) = resolved.place { return ty; } - constraint_keys.extend(keys); } } value_type .member(db, &attr.id) - .map_type(|ty| self.narrow(attribute, ty, constraint_keys)) + .map_type(|ty| self.narrow_expr_with_applicable_constraints(attribute, ty, &constraint_keys)) .unwrap_with_diagnostic(|lookup_error| match lookup_error { LookupError::Unbound(_) => { let report_unresolved_attribute = self.is_reachable(attribute); @@ -7698,6 +7692,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // If `value` is a valid reference, we attempt type narrowing by assignment. if !value_ty.is_unknown() { if let Ok(expr) = PlaceExpr::try_from(subscript) { + let (place, keys) = + self.infer_place_load(&expr, ast::ExprRef::Subscript(subscript)); + constraint_keys.extend(keys); // Type narrowing based on assignment to a subscript expression is generally // unsound, because arbitrary `__getitem__`/`__setitem__` methods on a class do not // necessarily guarantee that the passed-in value for `__setitem__` is stored and @@ -7726,13 +7723,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .zip(safe_mutable_class.generic_origin(db)) .is_some_and(|(l, r)| l == r) }) { - let (place, keys) = - self.infer_place_load(&expr, ast::ExprRef::Subscript(subscript)); if let Place::Type(ty, Boundness::Bound) = place.place { self.infer_expression(slice); return ty; } - constraint_keys.extend(keys); } } } @@ -7763,7 +7757,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let slice_ty = self.infer_expression(slice); let result_ty = self.infer_subscript_expression_types(value, value_ty, slice_ty); - self.narrow(subscript, result_ty, constraint_keys) + self.narrow_expr_with_applicable_constraints(subscript, result_ty, &constraint_keys) } fn infer_explicit_class_specialization( From 72a97f6deefa002d81b0a66c5e221819e812a783 Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Mon, 9 Jun 2025 15:41:10 +0900 Subject: [PATCH 09/17] refactor --- .../mdtest/narrow/{composite_target.md => complex_target.md} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename crates/ty_python_semantic/resources/mdtest/narrow/{composite_target.md => complex_target.md} (96%) diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/composite_target.md b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md similarity index 96% rename from crates/ty_python_semantic/resources/mdtest/narrow/composite_target.md rename to crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md index 140787dcf19dd8..cb41876c36e520 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/composite_target.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md @@ -1,4 +1,4 @@ -# Narrowing for composite targets +# Narrowing for complex targets ## Member access @@ -56,7 +56,7 @@ def _(d: dict[str, str | None]): reveal_type(d["b"]) # revealed: str | None ``` -## Complex target +## Member access and subscript ```py class C: From 20b76d04e1705d7554f57f5663ad1554550ec050 Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Mon, 9 Jun 2025 22:24:11 +0900 Subject: [PATCH 10/17] Move the tornado package to bad.txt --- crates/ty_python_semantic/resources/primer/bad.txt | 1 + crates/ty_python_semantic/resources/primer/good.txt | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/ty_python_semantic/resources/primer/bad.txt b/crates/ty_python_semantic/resources/primer/bad.txt index 41010335f3a01e..18dce54689abfe 100644 --- a/crates/ty_python_semantic/resources/primer/bad.txt +++ b/crates/ty_python_semantic/resources/primer/bad.txt @@ -18,4 +18,5 @@ setuptools # vendors packaging, see above spack # slow, success, but mypy-primer hangs processing the output spark # too many iterations steam.py # hangs (single threaded) +tornado # bad use-def map (https://github.com/astral-sh/ty/issues/365) xarray # too many iterations diff --git a/crates/ty_python_semantic/resources/primer/good.txt b/crates/ty_python_semantic/resources/primer/good.txt index ea322176ea05bb..ee24a584cce636 100644 --- a/crates/ty_python_semantic/resources/primer/good.txt +++ b/crates/ty_python_semantic/resources/primer/good.txt @@ -109,7 +109,6 @@ stone strawberry streamlit sympy -tornado trio twine typeshed-stats From b6a764717549209e6a83ba5ca7ff819658d2af7f Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama <45118249+mtshiba@users.noreply.github.com> Date: Fri, 13 Jun 2025 00:03:08 +0900 Subject: [PATCH 11/17] Update crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md Co-authored-by: David Peter --- .../resources/mdtest/narrow/complex_target.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md index cb41876c36e520..14111299d0d9ac 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md @@ -61,7 +61,7 @@ def _(d: dict[str, str | None]): ```py class C: def __init__(self): - self.x: tuple[int | None, int | None] = () + self.x: tuple[int | None, int | None] = (None, None) class D: def __init__(self): From 1a69abc4d260b9378ecc2eb161814ac391f43a92 Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Fri, 13 Jun 2025 00:10:07 +0900 Subject: [PATCH 12/17] refactor: `expr` -> `place` --- crates/ty_python_semantic/src/types/infer.rs | 48 ++++++++++---------- 1 file changed, 25 insertions(+), 23 deletions(-) diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index d25e3e66aa8cb7..e913bbb0fcfd8a 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -1502,13 +1502,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let global_use_def_map = self.index.use_def_map(FileScopeId::global()); let place_id = binding.place(self.db()); - let expr = place_table.place_expr(place_id); + let place = place_table.place_expr(place_id); let skip_non_global_scopes = self.skip_non_global_scopes(file_scope_id, place_id); let declarations = if skip_non_global_scopes { match self .index .place_table(FileScopeId::global()) - .place_id_by_expr(&expr.expr) + .place_id_by_expr(&place.expr) { Some(id) => global_use_def_map.public_declarations(id), // This case is a syntax error (load before global declaration) but ignore that here @@ -1519,18 +1519,20 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { }; let declared_ty = place_from_declarations(self.db(), declarations) - .and_then(|place| { - Ok(if matches!(place.place, Place::Type(_, Boundness::Bound)) { - place - } else if skip_non_global_scopes - || self.scope().file_scope_id(self.db()).is_global() - { - let module_type_declarations = - module_type_implicit_global_declaration(self.db(), &expr.expr)?; - place.or_fall_back_to(self.db(), || module_type_declarations) - } else { - place - }) + .and_then(|place_and_quals| { + Ok( + if matches!(place_and_quals.place, Place::Type(_, Boundness::Bound)) { + place_and_quals + } else if skip_non_global_scopes + || self.scope().file_scope_id(self.db()).is_global() + { + let module_type_declarations = + module_type_implicit_global_declaration(self.db(), &place.expr)?; + place_and_quals.or_fall_back_to(self.db(), || module_type_declarations) + } else { + place_and_quals + }, + ) }) .map( |PlaceAndQualifiers { @@ -1568,10 +1570,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { ) .unwrap_or_else(|(ty, conflicting)| { // TODO point out the conflicting declarations in the diagnostic? - let expr = place_table.place_expr(binding.place(db)); + let place = place_table.place_expr(binding.place(db)); if let Some(builder) = self.context.report_lint(&CONFLICTING_DECLARATIONS, node) { builder.into_diagnostic(format_args!( - "Conflicting declared types for `{expr}`: {}", + "Conflicting declared types for `{place}`: {}", conflicting.display(db) )); } @@ -1616,9 +1618,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // Fallback to bindings declared on `types.ModuleType` if it's a global symbol let scope = self.scope().file_scope_id(self.db()); let place_table = self.index.place_table(scope); - let expr = place_table.place_expr(declaration.place(self.db())); - if scope.is_global() && expr.is_name() { - module_type_implicit_global_symbol(self.db(), expr.expect_name()) + let place = place_table.place_expr(declaration.place(self.db())); + if scope.is_global() && place.is_name() { + module_type_implicit_global_symbol(self.db(), place.expect_name()) } else { Place::Unbound.into() } @@ -1669,9 +1671,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let file_scope_id = self.scope().file_scope_id(self.db()); if file_scope_id.is_global() { let place_table = self.index.place_table(file_scope_id); - let expr = place_table.place_expr(definition.place(self.db())); + let place = place_table.place_expr(definition.place(self.db())); if let Some(module_type_implicit_declaration) = - module_type_implicit_global_declaration(self.db(), &expr.expr) + module_type_implicit_global_declaration(self.db(), &place.expr) .ok() .and_then(|place| place.place.ignore_possibly_unbound()) { @@ -1683,11 +1685,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { self.context.report_lint(&INVALID_DECLARATION, node) { let mut diagnostic = builder.into_diagnostic(format_args!( - "Cannot shadow implicit global attribute `{expr}` with declaration of type `{}`", + "Cannot shadow implicit global attribute `{place}` with declaration of type `{}`", declared_type.display(self.db()) )); diagnostic.info(format_args!("The global symbol `{}` must always have a type assignable to `{}`", - expr, + place, module_type_implicit_declaration.display(self.db()) )); } From 325230cc986b37b1aba1f95c0caaa5c8615b2046 Mon Sep 17 00:00:00 2001 From: David Peter Date: Tue, 10 Jun 2025 13:06:04 +0200 Subject: [PATCH 13/17] Add tests that combine narrowing and assignment to attributes --- .../resources/mdtest/narrow/complex_target.md | 56 +++++++++++++++++-- 1 file changed, 52 insertions(+), 4 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md index 14111299d0d9ac..fea8053cc10c8c 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md @@ -1,6 +1,54 @@ -# Narrowing for complex targets +# Narrowing for complex targets (attribute expressions, subscripts) -## Member access +We support type narrowing for attributes and subscripts. + +## Basic attribute narrowing + +```py +class C: + x: int | None = None + +c = C() + +reveal_type(c.x) # revealed: int | None + +if c.x is not None: + reveal_type(c.x) # revealed: int +else: + reveal_type(c.x) # revealed: None +``` + +Narrowing can be "reset" by assigning to the attribute: + +```py +c = C() + +if c.x is None: + reveal_type(c.x) # revealed: None + c.x = 1 + reveal_type(c.x) # revealed: Literal[1] + c.x = None + reveal_type(c.x) # revealed: None + +# TODO: this should be int | None +reveal_type(c.x) # revealed: int +``` + +Narrowing can also be "reset" by assigning to the object: + +```py +c = C() + +if c.x is None: + reveal_type(c.x) # revealed: None + c = C() + reveal_type(c.x) # revealed: int | None + +# TODO: this should be int | None +reveal_type(c.x) # revealed: int +``` + +## Attribute narrowing with intermediate scopes ```py class C: @@ -23,7 +71,7 @@ def _(): reveal_type(c.x) # revealed: (Unknown & ~None) | int ``` -## Subscript +## Subscript narrowing ### Number subscript @@ -56,7 +104,7 @@ def _(d: dict[str, str | None]): reveal_type(d["b"]) # revealed: str | None ``` -## Member access and subscript +## Combined attribute and subscript narrowing ```py class C: From d8ca65ee70d742d07abdb250482559ae009dfd9e Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Fri, 13 Jun 2025 23:57:52 +0900 Subject: [PATCH 14/17] Update complex_target.md --- .../resources/mdtest/narrow/complex_target.md | 69 ++++++++++++++++++- 1 file changed, 67 insertions(+), 2 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md index fea8053cc10c8c..dc2e25a5591949 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md @@ -2,7 +2,9 @@ We support type narrowing for attributes and subscripts. -## Basic attribute narrowing +## Attribute narrowing + +### Basic ```py class C: @@ -48,7 +50,70 @@ if c.x is None: reveal_type(c.x) # revealed: int ``` -## Attribute narrowing with intermediate scopes +### Multiple predicates + +```py +class C: + value: str | None + +def foo(c: C): + if c.value and len(c.value): + reveal_type(c.value) # revealed: str & ~AlwaysFalsy + + # error: [invalid-argument-type] "Argument to function `len` is incorrect: Expected `Sized`, found `str | None`" + if len(c.value) and c.value: + reveal_type(c.value) # revealed: str & ~AlwaysFalsy + + if c.value is None or not len(c.value): + reveal_type(c.value) # revealed: str | None + else: # c.value is not None and len(c.value) + # TODO: should be # `str & ~AlwaysFalsy` + reveal_type(c.value) # revealed: str +``` + +### Generic class + +```toml +[environment] +python-version = "3.12" +``` + +```py +class C[T]: + x: T + y: T + + def __init__(self, x: T): + self.x = x + self.y = x + +def f(a: int | None): + c = C(a) + reveal_type(c.x) # revealed: int | None + reveal_type(c.y) # revealed: int | None + if c.x is not None: + reveal_type(c.x) # revealed: int + # In this case, it may seem like we can narrow it down to `int`, + # but different values ​​may be reassigned to `x` and `y` in another place. + reveal_type(c.y) # revealed: int | None + +def g[T](c: C[T]): + reveal_type(c.x) # revealed: T + reveal_type(c.y) # revealed: T + reveal_type(c) # revealed: C[T] + + if isinstance(c.x, int): + reveal_type(c.x) # revealed: T & int + reveal_type(c.y) # revealed: T + reveal_type(c) # revealed: C[T] + if isinstance(c.x, int) and isinstance(c.y, int): + reveal_type(c.x) # revealed: T & int + reveal_type(c.y) # revealed: T & int + # TODO: Probably better if inferred as `C[T & int]` (mypy and pyright don't support this) + reveal_type(c) # revealed: C[T] +``` + +### With intermediate scopes ```py class C: From 991f437b96324d52933f6954f1571120cc45ff6d Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Sun, 15 Jun 2025 18:29:51 +0900 Subject: [PATCH 15/17] =?UTF-8?q?Check=20attribute/subscript=20values=20?= =?UTF-8?q?=E2=80=8B=E2=80=8Bat=20store=20time,=20not=20at=20load=20time?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../mdtest/narrow/conditionals/nested.md | 3 +- crates/ty_python_semantic/src/types/infer.rs | 96 +++++++++++-------- 2 files changed, 58 insertions(+), 41 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md index bd73933b6362f1..8db508f91614f5 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md @@ -134,8 +134,7 @@ class _: class _3: reveal_type(a) # revealed: A - # TODO: should be `D | None` - reveal_type(a.b.c1.d) # revealed: D + reveal_type(a.b.c1.d) # revealed: D | None a.b.c1 = C() a.b.c1.d = D() diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index f8c23183d19316..d321248151d7b3 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -741,6 +741,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .expression_type(expr.scoped_expression_id(self.db(), self.scope())) } + fn try_expression_type(&self, expr: &ast::Expr) -> Option> { + self.types + .try_expression_type(expr.scoped_expression_id(self.db(), self.scope())) + } + /// Get the type of an expression from any scope in the same file. /// /// If the expression is in the current scope, and we are inferring the entire scope, just look @@ -1587,6 +1592,54 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // allow declarations to override inference in case of invalid assignment bound_ty = declared_ty; } + // In the following cases, the bound type may not be the same as the RHS value type. + if let AnyNodeRef::ExprAttribute(ast::ExprAttribute { value, attr, .. }) = node { + let value_ty = self + .try_expression_type(value) + .unwrap_or_else(|| self.infer_maybe_standalone_expression(value)); + // If the member is a data descriptor, the RHS value may differ from the value actually assigned. + if value_ty + .class_member(db, attr.id.clone()) + .place + .ignore_possibly_unbound() + .is_some_and(|ty| ty.may_be_data_descriptor(db)) + { + bound_ty = declared_ty; + } + } else if let AnyNodeRef::ExprSubscript(ast::ExprSubscript { value, .. }) = node { + let value_ty = self + .try_expression_type(value) + .unwrap_or_else(|| self.infer_expression(value)); + // Arbitrary `__getitem__`/`__setitem__` methods on a class do not + // necessarily guarantee that the passed-in value for `__setitem__` is stored and + // can be retrieved unmodified via `__getitem__`. Therefore, we currently only + // perform assignment-based narrowing on a few built-in classes (`list`, `dict`, + // `bytesarray`, `TypedDict` and `collections` types) where we are confident that + // this kind of narrowing can be performed soundly. This is the same approach as + // pyright. TODO: Other standard library classes may also be considered safe. Also, + // subclasses of these safe classes that do not override `__getitem__/__setitem__` + // may be considered safe. + let safe_mutable_classes = [ + KnownClass::List.to_instance(db), + KnownClass::Dict.to_instance(db), + KnownClass::Bytearray.to_instance(db), + KnownClass::DefaultDict.to_instance(db), + SpecialFormType::ChainMap.instance_fallback(db), + SpecialFormType::Counter.instance_fallback(db), + SpecialFormType::Deque.instance_fallback(db), + SpecialFormType::OrderedDict.instance_fallback(db), + SpecialFormType::TypedDict.instance_fallback(db), + ]; + if safe_mutable_classes.iter().all(|safe_mutable_class| { + !value_ty.is_equivalent_to(db, *safe_mutable_class) + && value_ty + .generic_origin(db) + .zip(safe_mutable_class.generic_origin(db)) + .is_none_or(|(l, r)| l != r) + }) { + bound_ty = declared_ty; + } + } self.types.bindings.insert(binding, bound_ty); } @@ -6191,15 +6244,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let (resolved, keys) = self.infer_place_load(&place_expr, ast::ExprRef::Attribute(attribute)); constraint_keys.extend(keys); - let member = value_type.class_member(db, attr.id.clone()); - // If the member is a data descriptor, the value most recently assigned - // to the attribute may not necessarily be obtained here. - if member - .place - .ignore_possibly_unbound() - .is_none_or(|ty| !ty.may_be_data_descriptor(db)) - { - if let Place::Type(ty, Boundness::Bound) = resolved.place { + if let Place::Type(ty, Boundness::Bound) = resolved.place { + if !ty.is_unknown() { return ty; } } @@ -7702,7 +7748,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { slice, ctx: _, } = subscript; - let db = self.db(); let value_ty = self.infer_expression(value); let mut constraint_keys = vec![]; @@ -7712,35 +7757,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let (place, keys) = self.infer_place_load(&expr, ast::ExprRef::Subscript(subscript)); constraint_keys.extend(keys); - // Type narrowing based on assignment to a subscript expression is generally - // unsound, because arbitrary `__getitem__`/`__setitem__` methods on a class do not - // necessarily guarantee that the passed-in value for `__setitem__` is stored and - // can be retrieved unmodified via `__getitem__`. Therefore, we currently only - // perform assignment-based narrowing on a few built-in classes (`list`, `dict`, - // `bytesarray`, `TypedDict` and `collections` types) where we are confident that - // this kind of narrowing can be performed soundly. This is the same approach as - // pyright. TODO: Other standard library classes may also be considered safe. Also, - // subclasses of these safe classes that do not override `__getitem__/__setitem__` - // may be considered safe. - let safe_mutable_classes = [ - KnownClass::List.to_instance(db), - KnownClass::Dict.to_instance(db), - KnownClass::Bytearray.to_instance(db), - KnownClass::DefaultDict.to_instance(db), - SpecialFormType::ChainMap.instance_fallback(db), - SpecialFormType::Counter.instance_fallback(db), - SpecialFormType::Deque.instance_fallback(db), - SpecialFormType::OrderedDict.instance_fallback(db), - SpecialFormType::TypedDict.instance_fallback(db), - ]; - if safe_mutable_classes.iter().any(|safe_mutable_class| { - value_ty.is_equivalent_to(db, *safe_mutable_class) - || value_ty - .generic_origin(db) - .zip(safe_mutable_class.generic_origin(db)) - .is_some_and(|(l, r)| l == r) - }) { - if let Place::Type(ty, Boundness::Bound) = place.place { + if let Place::Type(ty, Boundness::Bound) = place.place { + if !ty.is_unknown() { self.infer_expression(slice); return ty; } From 1292a5a550eba6e408a7ae12a3b2a1a86f4b21ed Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Mon, 16 Jun 2025 15:07:41 +0900 Subject: [PATCH 16/17] Fix the combination of conditional & assignment based narrowing to work correctly --- .../resources/mdtest/narrow/complex_target.md | 31 ++++- .../ty_python_semantic/src/semantic_index.rs | 4 +- .../src/semantic_index/builder.rs | 9 ++ .../semantic_index/narrowing_constraints.rs | 2 + .../src/semantic_index/use_def.rs | 43 +++++-- crates/ty_python_semantic/src/types/infer.rs | 110 ++++++++++++++---- 6 files changed, 160 insertions(+), 39 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md index dc2e25a5591949..bf08e91fa94bd0 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md @@ -18,6 +18,31 @@ if c.x is not None: reveal_type(c.x) # revealed: int else: reveal_type(c.x) # revealed: None + +if c.x is not None: + c.x = None + +reveal_type(c.x) # revealed: None + +c = C() + +if c.x is None: + c.x = 1 + +reveal_type(c.x) # revealed: int + +class _: + reveal_type(c.x) # revealed: int + +c = C() + +class _: + if c.x is None: + c.x = 1 + reveal_type(c.x) # revealed: int + +# TODO: should be `int` +reveal_type(c.x) # revealed: int | None ``` Narrowing can be "reset" by assigning to the attribute: @@ -32,8 +57,7 @@ if c.x is None: c.x = None reveal_type(c.x) # revealed: None -# TODO: this should be int | None -reveal_type(c.x) # revealed: int +reveal_type(c.x) # revealed: int | None ``` Narrowing can also be "reset" by assigning to the object: @@ -46,8 +70,7 @@ if c.x is None: c = C() reveal_type(c.x) # revealed: int | None -# TODO: this should be int | None -reveal_type(c.x) # revealed: int +reveal_type(c.x) # revealed: int | None ``` ### Multiple predicates diff --git a/crates/ty_python_semantic/src/semantic_index.rs b/crates/ty_python_semantic/src/semantic_index.rs index 0b71130d1eb08c..719c80f661df31 100644 --- a/crates/ty_python_semantic/src/semantic_index.rs +++ b/crates/ty_python_semantic/src/semantic_index.rs @@ -37,8 +37,8 @@ mod use_def; mod visibility_constraints; pub(crate) use self::use_def::{ - BindingWithConstraints, BindingWithConstraintsIterator, DeclarationWithConstraint, - DeclarationsIterator, + ApplicableConstraints, BindingWithConstraints, BindingWithConstraintsIterator, + DeclarationWithConstraint, DeclarationsIterator, }; type PlaceSet = hashbrown::HashMap; diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index baafb07f343fa7..d1ed290b9fb94c 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -296,6 +296,15 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { // If the scope that we just popped off is an eager scope, we need to "lock" our view of // which bindings reach each of the uses in the scope. Loop through each enclosing scope, // looking for any that bind each place. + // TODO: Bindings in eager nested scopes also need to be recorded. For example: + // ```python + // class C: + // x: int | None = None + // c = C() + // class _: + // c.x = 1 + // reveal_type(c.x) # revealed: Literal[1] + // ``` for enclosing_scope_info in self.scope_stack.iter().rev() { let enclosing_scope_id = enclosing_scope_info.file_scope_id; let enclosing_scope_kind = self.scopes[enclosing_scope_id].kind(); diff --git a/crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs b/crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs index fa6280ead67cee..48297e1da6c4c3 100644 --- a/crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs +++ b/crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs @@ -30,6 +30,7 @@ use crate::list::{List, ListBuilder, ListSetReverseIterator, ListStorage}; use crate::semantic_index::ast_ids::ScopedUseId; +use crate::semantic_index::place::FileScopeId; use crate::semantic_index::predicate::ScopedPredicateId; /// A narrowing constraint associated with a live binding. @@ -42,6 +43,7 @@ pub(crate) type ScopedNarrowingConstraint = List { scope_start_visibility: ScopedVisibilityConstraintId, } +pub(crate) enum ApplicableConstraints<'map, 'db> { + UnboundBinding(ConstraintsIterator<'map, 'db>), + ConstrainedBindings(BindingWithConstraintsIterator<'map, 'db>), +} + impl<'db> UseDefMap<'db> { pub(crate) fn bindings_at_use( &self, @@ -359,19 +366,33 @@ impl<'db> UseDefMap<'db> { self.bindings_iterator(&self.bindings_by_use[use_id]) } - pub(crate) fn narrowing_constraints_at_use( + pub(crate) fn applicable_constraints( &self, constraint_key: ConstraintKey, - ) -> ConstraintsIterator<'_, 'db> { - let constraint = match constraint_key { - ConstraintKey::NarrowingConstraint(constraint) => constraint, + enclosing_scope: FileScopeId, + expr: &PlaceExpr, + index: &'db SemanticIndex, + ) -> ApplicableConstraints<'_, 'db> { + match constraint_key { + ConstraintKey::NarrowingConstraint(constraint) => { + ApplicableConstraints::UnboundBinding(ConstraintsIterator { + predicates: &self.predicates, + constraint_ids: self.narrowing_constraints.iter_predicates(constraint), + }) + } + ConstraintKey::EagerNestedScope(nested_scope) => { + let EagerSnapshotResult::FoundBindings(bindings) = + index.eager_snapshot(enclosing_scope, expr, nested_scope) + else { + unreachable!( + "The result of `SemanticIndex::eager_snapshot` must be `FoundBindings`" + ) + }; + ApplicableConstraints::ConstrainedBindings(bindings) + } ConstraintKey::UseId(use_id) => { - self.bindings_by_use[use_id].unbound_narrowing_constraint() + ApplicableConstraints::ConstrainedBindings(self.bindings_at_use(use_id)) } - }; - ConstraintsIterator { - predicates: &self.predicates, - constraint_ids: self.narrowing_constraints.iter_predicates(constraint), } } diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index d321248151d7b3..9658ef4f13be02 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -60,7 +60,7 @@ use crate::semantic_index::ast_ids::{ }; use crate::semantic_index::definition::{ AnnotatedAssignmentDefinitionKind, AssignmentDefinitionKind, ComprehensionDefinitionKind, - Definition, DefinitionKind, DefinitionNodeKey, ExceptHandlerDefinitionKind, + Definition, DefinitionKind, DefinitionNodeKey, DefinitionState, ExceptHandlerDefinitionKind, ForStmtDefinitionKind, TargetKind, WithItemDefinitionKind, }; use crate::semantic_index::expression::{Expression, ExpressionKind}; @@ -68,7 +68,9 @@ use crate::semantic_index::narrowing_constraints::ConstraintKey; use crate::semantic_index::place::{ FileScopeId, NodeWithScopeKind, NodeWithScopeRef, PlaceExpr, ScopeId, ScopeKind, ScopedPlaceId, }; -use crate::semantic_index::{EagerSnapshotResult, SemanticIndex, semantic_index}; +use crate::semantic_index::{ + ApplicableConstraints, EagerSnapshotResult, SemanticIndex, semantic_index, +}; use crate::types::call::{ Argument, Binding, Bindings, CallArgumentTypes, CallArguments, CallError, }; @@ -5837,11 +5839,69 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let db = self.db(); for (enclosing_scope_file_id, constraint_key) in constraint_keys { let use_def = self.index.use_def_map(*enclosing_scope_file_id); - let constraints = use_def.narrowing_constraints_at_use(*constraint_key); let place_table = self.index.place_table(*enclosing_scope_file_id); let place = place_table.place_id_by_expr(expr).unwrap(); - ty = constraints.narrow(db, ty, place); + match use_def.applicable_constraints( + *constraint_key, + *enclosing_scope_file_id, + expr, + self.index, + ) { + ApplicableConstraints::UnboundBinding(constraint) => { + ty = constraint.narrow(db, ty, place); + } + // Performs narrowing based on constrained bindings. + // This handling must be performed even if narrowing is attempted and failed using `infer_place_load`. + // The result of `infer_place_load` can be applied as is only when its boundness is `Bound`. + // For example, this handling is required in the following case: + // ```python + // class C: + // x: int | None = None + // c = C() + // # c.x: int | None = + // if c.x is None: + // c.x = 1 + // # else: c.x: int = + // # `c.x` is not definitely bound here + // reveal_type(c.x) # revealed: int + // ``` + ApplicableConstraints::ConstrainedBindings(bindings) => { + let visibility_constraints = bindings.visibility_constraints; + let predicates = bindings.predicates; + let mut union = UnionBuilder::new(db); + for binding in bindings { + let static_visibility = visibility_constraints.evaluate( + db, + predicates, + binding.visibility_constraint, + ); + if static_visibility.is_always_false() { + continue; + } + match binding.binding { + DefinitionState::Defined(definition) => { + let binding_ty = binding_type(db, definition); + union = union.add( + binding.narrowing_constraint.narrow(db, binding_ty, place), + ); + } + DefinitionState::Undefined | DefinitionState::Deleted => { + union = + union.add(binding.narrowing_constraint.narrow(db, ty, place)); + } + } + } + // If there are no visible bindings, the union becomes `Never`. + // Since an unbound binding is recorded even for an undefined place, + // this can only happen if the code is unreachable + // and therefore it is correct to set the result to `Never`. + let union = union.build(); + if !union.is_unknown() && union.is_assignable_to(db, ty) { + ty = union; + } + } + } } ty } @@ -6053,15 +6113,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { { continue; } - return place_from_bindings(db, bindings) - .map_type(|ty| { - self.narrow_place_with_applicable_constraints( - expr, - ty, - &constraint_keys, - ) - }) - .into(); + let place = place_from_bindings(db, bindings).map_type(|ty| { + self.narrow_place_with_applicable_constraints( + expr, + ty, + &constraint_keys, + ) + }); + constraint_keys.push(( + enclosing_scope_file_id, + ConstraintKey::EagerNestedScope(file_scope_id), + )); + return place.into(); } // There are no visible bindings / constraint here. // Don't fall back to non-eager place resolution. @@ -6125,15 +6188,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { )); } EagerSnapshotResult::FoundBindings(bindings) => { - return place_from_bindings(db, bindings) - .map_type(|ty| { - self.narrow_place_with_applicable_constraints( - expr, - ty, - &constraint_keys, - ) - }) - .into(); + let place = place_from_bindings(db, bindings).map_type(|ty| { + self.narrow_place_with_applicable_constraints( + expr, + ty, + &constraint_keys, + ) + }); + constraint_keys.push(( + FileScopeId::global(), + ConstraintKey::EagerNestedScope(file_scope_id), + )); + return place.into(); } // There are no visible bindings / constraint here. EagerSnapshotResult::NotFound => { From 96ed63bc0ca84126402a19e6c5dc3ce61e092536 Mon Sep 17 00:00:00 2001 From: Shunsuke Shibayama Date: Mon, 16 Jun 2025 23:59:23 +0900 Subject: [PATCH 17/17] Perform default type inference even if we can obtain the type based on the assignments --- .../resources/mdtest/attributes.md | 3 ++- .../resources/mdtest/narrow/assignment.md | 6 +++++ .../resources/mdtest/narrow/complex_target.md | 15 ++++++++++++ .../mdtest/narrow/conditionals/nested.md | 3 ++- crates/ty_python_semantic/src/types/infer.rs | 24 +++++++++++-------- 5 files changed, 39 insertions(+), 12 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/attributes.md b/crates/ty_python_semantic/resources/mdtest/attributes.md index 71288e51093439..53643da8ecdb09 100644 --- a/crates/ty_python_semantic/resources/mdtest/attributes.md +++ b/crates/ty_python_semantic/resources/mdtest/attributes.md @@ -751,7 +751,8 @@ reveal_type(C.pure_class_variable) # revealed: Unknown # and the assignment is properly attributed to the class method. # error: [invalid-attribute-access] "Cannot assign to instance attribute `pure_class_variable` from the class object ``" C.pure_class_variable = "overwritten on class" - +# TODO: should be no error +# error: [unresolved-attribute] "Attribute `pure_class_variable` can only be accessed on instances, not on the class object `` itself." reveal_type(C.pure_class_variable) # revealed: Literal["overwritten on class"] c_instance = C() diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md b/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md index 73d676a2a371b6..d5a59ad275ceff 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/assignment.md @@ -87,6 +87,12 @@ class _: reveal_type(a.y) # revealed: Unknown | None reveal_type(a.z) # revealed: Unknown | None +a = A() +# error: [unresolved-attribute] +a.dynamically_added = 0 +# error: [unresolved-attribute] +reveal_type(a.dynamically_added) # revealed: Literal[0] + # error: [unresolved-reference] does.nt.exist = 0 # error: [unresolved-reference] diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md index bf08e91fa94bd0..c74f24a0aa3b41 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md @@ -7,6 +7,8 @@ We support type narrowing for attributes and subscripts. ### Basic ```py +from ty_extensions import Unknown + class C: x: int | None = None @@ -43,6 +45,19 @@ class _: # TODO: should be `int` reveal_type(c.x) # revealed: int | None + +class D: + x = None + +def unknown() -> Unknown: + return 1 + +d = D() +reveal_type(d.x) # revealed: Unknown | None +d.x = 1 +reveal_type(d.x) # revealed: Literal[1] +d.x = unknown() +reveal_type(d.x) # revealed: Unknown ``` Narrowing can be "reset" by assigning to the attribute: diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md index 8db508f91614f5..033ca89d3e4642 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md @@ -134,7 +134,8 @@ class _: class _3: reveal_type(a) # revealed: A - reveal_type(a.b.c1.d) # revealed: D | None + # TODO: should be `D | None` + reveal_type(a.b.c1.d) # revealed: Unknown a.b.c1 = C() a.b.c1.d = D() diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index c8f79485f24bff..ba87963d3ace8f 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -6038,7 +6038,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // this can only happen if the code is unreachable // and therefore it is correct to set the result to `Never`. let union = union.build(); - if !union.is_unknown() && union.is_assignable_to(db, ty) { + if union.is_assignable_to(db, ty) { ty = union; } } @@ -6449,18 +6449,17 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let db = self.db(); let mut constraint_keys = vec![]; + let mut assigned_type = None; if let Ok(place_expr) = PlaceExpr::try_from(attribute) { let (resolved, keys) = self.infer_place_load(&place_expr, ast::ExprRef::Attribute(attribute)); constraint_keys.extend(keys); if let Place::Type(ty, Boundness::Bound) = resolved.place { - if !ty.is_unknown() { - return ty; - } + assigned_type = Some(ty); } } - value_type + let resolved_type = value_type .member(db, &attr.id) .map_type(|ty| self.narrow_expr_with_applicable_constraints(attribute, ty, &constraint_keys)) .unwrap_with_diagnostic(|lookup_error| match lookup_error { @@ -6522,7 +6521,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { type_when_bound } - }).inner_type() + }) + .inner_type(); + // Even if we can obtain the attribute type based on the assignments, we still perform default type inference + // (to report errors). + assigned_type.unwrap_or(resolved_type) } fn infer_attribute_expression(&mut self, attribute: &ast::ExprAttribute) -> Type<'db> { @@ -7977,10 +7980,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { self.infer_place_load(&expr, ast::ExprRef::Subscript(subscript)); constraint_keys.extend(keys); if let Place::Type(ty, Boundness::Bound) = place.place { - if !ty.is_unknown() { - self.infer_expression(slice); - return ty; - } + // Even if we can obtain the subscript type based on the assignments, we still perform default type inference + // (to store the expression type and to report errors). + let slice_ty = self.infer_expression(slice); + self.infer_subscript_expression_types(value, value_ty, slice_ty); + return ty; } } }