[ty] Read inference and IDE inputs through typed AST getters
diff --git a/crates/ty_ide/src/symbols.rs b/crates/ty_ide/src/symbols.rs index 6d4b6fb..7b43b22 100644 --- a/crates/ty_ide/src/symbols.rs +++ b/crates/ty_ide/src/symbols.rs
@@ -1298,8 +1298,8 @@ } fn visit_stmt_impl(&mut self, stmt: Stmt<'db>) { - match stmt.fields() { - ast::StmtFields::FunctionDef(ast::StmtFunctionDefFields { node: func_def, .. }) => { + match stmt.classify() { + ast::StmtKind::FunctionDef(func_def) => { let kind = SymbolKind::function_kind( func_def.name().as_str(), self.iter_symbol_stack() @@ -1349,9 +1349,7 @@ self.pop_symbol(); } - ast::StmtFields::ClassDef(ast::StmtClassDefFields { - node: class_def, .. - }) => { + ast::StmtKind::ClassDef(class_def) => { let symbol = SymbolTree { parent: None, name: class_def.name().to_string(), @@ -1391,9 +1389,7 @@ // Restore the previous class scope state self.in_class = was_in_class; } - ast::StmtFields::TypeAlias(ast::StmtTypeAliasFields { - node: type_alias, .. - }) => { + ast::StmtKind::TypeAlias(type_alias) => { // Include assignments only when we're in global or class scope if self.in_function { return; @@ -1403,19 +1399,18 @@ }; self.add_name_symbol(stmt, name.id(), name.range(), SymbolKind::Variable); } - ast::StmtFields::Assign(ast::StmtAssignFields { node: assign, .. }) => { + ast::StmtKind::Assign(assign) => { self.add_all_assignment(assign.targets().first(), Some(assign.value())); source_order::walk_stmt(self, stmt); } - ast::StmtFields::AnnAssign(ast::StmtAnnAssignFields { - node: ann_assign, .. - }) => { + ast::StmtKind::AnnAssign(ann_assign) => { self.add_all_assignment(Some(ann_assign.target()), ann_assign.value()); source_order::walk_stmt(self, stmt); } - ast::StmtFields::AugAssign(ast::StmtAugAssignFields { - target, op, value, .. - }) => { + ast::StmtKind::AugAssign(assign) => { + let ast::StmtAugAssignFields { + target, op, value, .. + } = assign.fields(); if !target.is_name_expr() { self.visit_expr(target); } @@ -1444,7 +1439,7 @@ self.all_invalid = true; } } - ast::StmtFields::Expr(ast::StmtExprFields { node: expr, .. }) => { + ast::StmtKind::Expr(expr) => { source_order::walk_stmt(self, stmt); // We don't care about `__all__` unless we're @@ -1479,7 +1474,7 @@ self.all_invalid = true; } } - ast::StmtFields::Import(ast::StmtImportFields { node: import, .. }) => { + ast::StmtKind::Import(import) => { // We ignore any names introduced by imports // unless we're specifically looking for the // set of exported symbols. @@ -1495,9 +1490,7 @@ self.add_import_alias(AstImport::Import(import), alias); } } - ast::StmtFields::ImportFrom(ast::StmtImportFromFields { - node: import_from, .. - }) => { + ast::StmtKind::ImportFrom(import_from) => { // We ignore any names introduced by imports // unless we're specifically looking for the // set of exported symbols. @@ -1628,7 +1621,8 @@ /// Checks if the given expression is a name expression for `__all__`. fn is_dunder_all(expr: Expr<'_>) -> bool { - matches!(expr.fields(), ast::ExprFields::Name(ast::ExprNameFields { id, .. }) if id == "__all__") + expr.as_name_expr() + .is_some_and(|name| name.id() == "__all__") } fn is_recognized_all_assignment(stmt: Stmt<'_>, name: ExprName<'_>) -> bool {
diff --git a/crates/ty_python_core/src/lib.rs b/crates/ty_python_core/src/lib.rs index c3bf693..4066416 100644 --- a/crates/ty_python_core/src/lib.rs +++ b/crates/ty_python_core/src/lib.rs
@@ -1217,7 +1217,6 @@ files::{File, system_path_to_file}, parsed::ParsedModuleRef, }; - use ruff_python_ast as ast; use ruff_text_size::{Ranged, TextRange}; use test_case::test_case; @@ -2060,11 +2059,7 @@ let DefinitionKind::Assignment(assignment) = binding.kind(&db) else { panic!("should be an assignment definition") }; - let ast::ExprFields::NumberLiteral(ast::ExprNumberLiteralFields { - value: ast::Number::Int(num), - .. - }) = assignment.value(&module).fields() - else { + let Some(num) = assignment.value(&module).as_int_literal() else { panic!("should be a number literal") }; assert_eq!(*num, 1);
diff --git a/crates/ty_python_core/src/member.rs b/crates/ty_python_core/src/member.rs index dc396db..0399f81 100644 --- a/crates/ty_python_core/src/member.rs +++ b/crates/ty_python_core/src/member.rs
@@ -304,59 +304,39 @@ } fn subscript_part(subscript_slice: Expr<'_>) -> Option<(SegmentKind, MemberPathPart<'_>)> { - match subscript_slice.fields() { + match subscript_slice.classify() { // Handle integer subscripts, like `x[0]`. - ast::ExprFields::NumberLiteral(ast::ExprNumberLiteralFields { - value: Number::Int(index), - .. - }) => Some(( - SegmentKind::IntSubscript, - MemberPathPart::Owned(format_char!("{index}")), - )), - // Handle negative integer subscripts, like `x[-1]`. - ast::ExprFields::UnaryOp(ast::ExprUnaryOpFields { - op: UnaryOp::USub, - operand, - .. - }) => match operand.fields() { - ast::ExprFields::NumberLiteral(ast::ExprNumberLiteralFields { - value: Number::Int(index), - .. - }) => Some(( - SegmentKind::IntSubscript, - MemberPathPart::Owned(format_char!("-{index}")), - )), - _ => None, - }, - // Handle positive integer subscripts with explicit plus, like `x[+1]`. - ast::ExprFields::UnaryOp(ast::ExprUnaryOpFields { - op: UnaryOp::UAdd, - operand, - .. - }) => match operand.fields() { - ast::ExprFields::NumberLiteral(ast::ExprNumberLiteralFields { - value: Number::Int(index), - .. - }) => Some(( + ast::ExprKind::NumberLiteral(number) => { + let Number::Int(index) = number.value() else { + return None; + }; + Some(( SegmentKind::IntSubscript, MemberPathPart::Owned(format_char!("{index}")), - )), - _ => None, - }, + )) + } + // Handle signed integer subscripts, like `x[-1]` and `x[+1]`. + ast::ExprKind::UnaryOp(unary) => { + let index = unary.operand().as_int_literal()?; + let index = match unary.op() { + UnaryOp::USub => format_char!("-{index}"), + UnaryOp::UAdd => format_char!("{index}"), + UnaryOp::Not | UnaryOp::Invert => return None, + }; + Some((SegmentKind::IntSubscript, MemberPathPart::Owned(index))) + } // Handle boolean subscripts, like `x[True]` or `x[False]`. // In Python, `True` and `False` are equivalent to `1` and `0` for indexing. - ast::ExprFields::BooleanLiteral(ast::ExprBooleanLiteralFields { value, .. }) => Some(( + ast::ExprKind::BooleanLiteral(boolean) => Some(( SegmentKind::IntSubscript, - MemberPathPart::Borrowed(if value { "1" } else { "0" }), + MemberPathPart::Borrowed(if boolean.value() { "1" } else { "0" }), )), - ast::ExprFields::StringLiteral(ast::ExprStringLiteralFields { - node: string, .. - }) => Some(( + ast::ExprKind::StringLiteral(string) => Some(( SegmentKind::StringSubscript, MemberPathPart::Borrowed(string.value().to_str()), )), // Handle bytes literal subscripts, like `x[b"key"]`. - ast::ExprFields::BytesLiteral(ast::ExprBytesLiteralFields { node: bytes, .. }) => { + ast::ExprKind::BytesLiteral(bytes) => { let bytes_vec: Vec<u8> = bytes.value().bytes().collect(); let text = String::from_utf8_lossy(&bytes_vec); Some((
diff --git a/crates/ty_python_semantic/src/dunder_all.rs b/crates/ty_python_semantic/src/dunder_all.rs index 94b0c71..e37713c 100644 --- a/crates/ty_python_semantic/src/dunder_all.rs +++ b/crates/ty_python_semantic/src/dunder_all.rs
@@ -453,7 +453,8 @@ /// Checks if the given expression is a name expression for `__all__`. fn is_dunder_all(expr: Expr<'_>) -> bool { - matches!(expr.fields(), ast::ExprFields::Name(ast::ExprNameFields { id, .. }) if id == "__all__") + expr.as_name_expr() + .is_some_and(|name| name.id() == "__all__") } /// Create and return a [`Name`] from the given expression, [`None`] if it is an invalid expression
diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 9c04bcf..ff91303 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs
@@ -1685,41 +1685,42 @@ /// normal attribute or subscript lookup on its receiver. fn fallback_member_declared_type(&mut self, node: SyntaxNode<'_>) -> Option<Type<'db>> { let db = self.db(); - if let ast::AnyFields::ExprAttribute(ast::ExprAttributeFields { value, attr, .. }) = - node.fields() - { - let value_type = self.try_expression_type(value).unwrap_or_else(|| { - self.infer_maybe_standalone_expression(value, TypeContext::default()) - }); - if let Place::Defined(DefinedPlace { - ty, - definedness: Definedness::AlwaysDefined, - .. - }) = value_type - .member(db, self.program_environment(), attr.as_str()) - .place - { - // TODO: also consider qualifiers on the attribute - Some(ty) - } else { - None + match node.classify() { + ast::AnyKind::ExprAttribute(attribute) => { + let value = attribute.value(); + let value_type = self.try_expression_type(value).unwrap_or_else(|| { + self.infer_maybe_standalone_expression(value, TypeContext::default()) + }); + if let Place::Defined(DefinedPlace { + ty, + definedness: Definedness::AlwaysDefined, + .. + }) = value_type + .member(db, self.program_environment(), attribute.attr().as_str()) + .place + { + // TODO: also consider qualifiers on the attribute + Some(ty) + } else { + None + } } - } else if let ast::AnyFields::ExprSubscript(ast::ExprSubscriptFields { - node: subscript, - value, - slice, - ctx, - .. - }) = node.fields() - { - let value_ty = self.get_or_infer_expression(value, TypeContext::default()); - let slice_ty = self.get_or_infer_expression(slice, TypeContext::default()); - Some( - self.infer_subscript_expression_types(subscript, value_ty, slice_ty, ctx) + ast::AnyKind::ExprSubscript(subscript) => { + let value_ty = + self.get_or_infer_expression(subscript.value(), TypeContext::default()); + let slice_ty = + self.get_or_infer_expression(subscript.slice(), TypeContext::default()); + Some( + self.infer_subscript_expression_types( + subscript, + value_ty, + slice_ty, + subscript.ctx(), + ) .unwrap_or_else(|recovery_ty| recovery_ty), - ) - } else { - None + ) + } + _ => None, } } @@ -13328,9 +13329,9 @@ bound_ty = declared_ty; } // In the following cases, the bound type may not be the same as the RHS value type. - if let ast::AnyFields::ExprAttribute(ast::ExprAttributeFields { value, attr, .. }) = - self.node.fields() - { + if let Some(attribute) = self.node.cast::<ast::ExprAttribute>() { + let value = attribute.value(); + let attr = attribute.attr(); let value_ty = builder.try_expression_type(value).unwrap_or_else(|| { builder.infer_maybe_standalone_expression(value, TypeContext::default()) }); @@ -13381,9 +13382,8 @@ builder.discard_dict_key_assignments_for(self.binding); bound_ty = declared_ty; } - } else if let ast::AnyFields::ExprSubscript(ast::ExprSubscriptFields { value, .. }) = - self.node.fields() - { + } else if let Some(subscript) = self.node.cast::<ast::ExprSubscript>() { + let value = subscript.value(); let value_ty = builder .try_expression_type(value) .unwrap_or_else(|| builder.infer_expression(value, TypeContext::default()));
diff --git a/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs b/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs index faf600a..19972b1 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs
@@ -1687,12 +1687,10 @@ ); // Determine if this element unpacks a tuple: either `*expr` or `Unpack[expr]` let is_unpack = matches!(element.classify(), ExprKind::Starred(_)) - || matches!( - element.fields(), - ast::ExprFields::Subscript(ast::ExprSubscriptFields { value, .. }) - if self.expression_type(value) - == Type::SpecialForm(SpecialFormType::Unpack) - ); + || element.as_subscript_expr().is_some_and(|subscript| { + self.expression_type(subscript.value()) + == Type::SpecialForm(SpecialFormType::Unpack) + }); if is_unpack { let mut report_too_many_unpacked_tuples = || { @@ -1786,12 +1784,10 @@ return TupleType::homogeneous(db, env, Type::unknown()); } let single_element_is_unpack = slice.is_starred_expr() - || matches!( - slice.fields(), - ast::ExprFields::Subscript(ast::ExprSubscriptFields { value, .. }) - if self.expression_type(value) - == Type::SpecialForm(SpecialFormType::Unpack) - ); + || slice.as_subscript_expr().is_some_and(|subscript| { + self.expression_type(subscript.value()) + == Type::SpecialForm(SpecialFormType::Unpack) + }); if single_element_is_unpack { if let Some(inner_tuple) = single_element_ty.exact_tuple_instance_spec(self.db()) @@ -3210,13 +3206,7 @@ operand, .. }) if matches!(op, ast::UnaryOp::USub | ast::UnaryOp::UAdd) - && matches!( - operand.fields(), - ast::ExprFields::NumberLiteral(ast::ExprNumberLiteralFields { - value: ast::Number::Int(_), - .. - }) - ) => + && operand.as_int_literal().is_some() => { let ty = self.infer_unary_expression(unary); self.store_expression_type(parameters, ty);