From 239cbc6f33396615103f7da86246df330db46852 Mon Sep 17 00:00:00 2001 From: David Peter Date: Tue, 5 Nov 2024 20:25:45 +0100 Subject: [PATCH] [red-knot] Store starred-expression annotation types (#14106) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Store the expression type for annotations that are starred expressions (see [discussion here](https://github.com/astral-sh/ruff/pull/14091#discussion_r1828332857)) - Use `self.store_expression_type(…)` consistently throughout, as it makes sure that no double-insertion errors occur. closes #14115 ## Test Plan Added an invalid-syntax example to the corpus which leads to a panic on `main`. Also added a Markdown test with a valid-syntax example that would lead to a panic once we implement function parameter inference. --------- Co-authored-by: Alex Waygood --- .../mdtest/annotations/starred_expressions.md | 18 ++++++++ .../src/semantic_index/ast_ids.rs | 8 ++++ .../src/types/infer.rs | 41 ++++++++++--------- .../corpus/25_func_annotations_starred.py | 8 ++++ 4 files changed, 55 insertions(+), 20 deletions(-) create mode 100644 crates/red_knot_python_semantic/resources/mdtest/annotations/starred_expressions.md create mode 100644 crates/red_knot_workspace/resources/test/corpus/25_func_annotations_starred.py diff --git a/crates/red_knot_python_semantic/resources/mdtest/annotations/starred_expressions.md b/crates/red_knot_python_semantic/resources/mdtest/annotations/starred_expressions.md new file mode 100644 index 0000000000..a125afce48 --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/annotations/starred_expressions.md @@ -0,0 +1,18 @@ +# Starred expression annotations + +Type annotations for `*args` can be starred expressions themselves: + +```py +from typing_extensions import TypeVarTuple + +Ts = TypeVarTuple("Ts") + +def append_int(*args: *Ts) -> tuple[*Ts, int]: + # TODO: should show some representation of the variadic generic type + reveal_type(args) # revealed: @Todo + + return (*args, 1) + +# TODO should be tuple[Literal[True], Literal["a"], int] +reveal_type(append_int(True, "a")) # revealed: @Todo +``` diff --git a/crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs b/crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs index 4e620e71cd..61e5081b44 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs @@ -87,6 +87,14 @@ pub trait HasScopedAstId { fn scoped_ast_id(&self, db: &dyn Db, scope: ScopeId) -> Self::Id; } +impl HasScopedAstId for Box { + type Id = ::Id; + + fn scoped_ast_id(&self, db: &dyn Db, scope: ScopeId) -> Self::Id { + self.as_ref().scoped_ast_id(db, scope) + } +} + /// Uniquely identifies an [`ast::Expr`] in a [`crate::semantic_index::symbol::FileScopeId`]. #[newtype_index] pub struct ScopedExpressionId; diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index c4e0594a1f..1ca5b17313 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -1443,8 +1443,8 @@ impl<'db> TypeInferenceBuilder<'db> { TargetKind::Name => value_ty, }; + self.store_expression_type(name, target_ty); self.add_binding(name.into(), definition, target_ty); - self.types.expressions.insert(name_ast_id, target_ty); } fn infer_annotated_assignment_statement(&mut self, assignment: &ast::StmtAnnAssign) { @@ -1697,10 +1697,7 @@ impl<'db> TypeInferenceBuilder<'db> { .unwrap_with_diagnostic(iterable.into(), &mut self.diagnostics) }; - self.types.expressions.insert( - target.scoped_ast_id(self.db, self.scope()), - loop_var_value_ty, - ); + self.store_expression_type(target, loop_var_value_ty); self.add_binding(target.into(), definition, loop_var_value_ty); } @@ -2027,7 +2024,11 @@ impl<'db> TypeInferenceBuilder<'db> { ty } - fn store_expression_type(&mut self, expression: &ast::Expr, ty: Type<'db>) { + fn store_expression_type( + &mut self, + expression: &impl HasScopedAstId, + ty: Type<'db>, + ) { let expr_id = expression.scoped_ast_id(self.db, self.scope()); let previous = self.types.expressions.insert(expr_id, ty); assert_eq!(previous, None); @@ -3803,31 +3804,31 @@ impl<'db> TypeInferenceBuilder<'db> { impl<'db> TypeInferenceBuilder<'db> { fn infer_annotation_expression(&mut self, expression: &ast::Expr) -> Type<'db> { // https://typing.readthedocs.io/en/latest/spec/annotations.html#grammar-token-expression-grammar-annotation_expression - match expression { + let annotation_ty = match expression { // TODO: parse the expression and check whether it is a string annotation, since they // can be annotation expressions distinct from type expressions. // https://typing.readthedocs.io/en/latest/spec/annotations.html#string-annotations - ast::Expr::StringLiteral(_literal) => { - self.store_expression_type(expression, Type::Todo); - Type::Todo - } + ast::Expr::StringLiteral(_literal) => Type::Todo, // Annotation expressions also get special handling for `*args` and `**kwargs`. ast::Expr::Starred(starred) => self.infer_starred_expression(starred), // All other annotation expressions are (possibly) valid type expressions, so handle // them there instead. - type_expr => self.infer_type_expression(type_expr), - } + type_expr => self.infer_type_expression_no_store(type_expr), + }; + + self.store_expression_type(expression, annotation_ty); + annotation_ty } } /// Type expressions impl<'db> TypeInferenceBuilder<'db> { - fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> { + fn infer_type_expression_no_store(&mut self, expression: &ast::Expr) -> Type<'db> { // https://typing.readthedocs.io/en/latest/spec/annotations.html#grammar-token-expression-grammar-type_expression - let ty = match expression { + match expression { ast::Expr::Name(name) => { debug_assert_eq!(name.ctx, ast::ExprContext::Load); self.infer_name_expression(name).to_instance(self.db) @@ -3983,12 +3984,12 @@ impl<'db> TypeInferenceBuilder<'db> { } ast::Expr::IpyEscapeCommand(_) => todo!("Implement Ipy escape command support"), - }; - - let expr_id = expression.scoped_ast_id(self.db, self.scope()); - let previous = self.types.expressions.insert(expr_id, ty); - assert!(previous.is_none()); + } + } + fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> { + let ty = self.infer_type_expression_no_store(expression); + self.store_expression_type(expression, ty); ty } diff --git a/crates/red_knot_workspace/resources/test/corpus/25_func_annotations_starred.py b/crates/red_knot_workspace/resources/test/corpus/25_func_annotations_starred.py new file mode 100644 index 0000000000..f48393b640 --- /dev/null +++ b/crates/red_knot_workspace/resources/test/corpus/25_func_annotations_starred.py @@ -0,0 +1,8 @@ +# Regression test for https://github.com/astral-sh/ruff/issues/14115 +# +# This is invalid syntax, but should not lead to a crash. + +def f() -> *int: ... + + +f()