Update type annotation parsing API to return Parsed (#11739)

## Summary This PR updates the return type of `parse_type_annotation` from `Expr` to `Parsed<ModExpression>`. This is to allow accessing the tokens for the parsed sub-expression in the follow-up PR. ## Test Plan `cargo insta test`
astral-sh · Jun 5, 2024 · eed6d78 · eed6d78
1 parent 8338db6
commit eed6d78
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 25 deletions.
diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs
@@ -30,13 +30,6 @@ use std::path::Path;
 
 use itertools::Itertools;
 use log::debug;
-use ruff_python_ast::{
-    self as ast, AnyParameterRef, Comprehension, ElifElseClause, ExceptHandler, Expr, ExprContext,
-    FStringElement, Keyword, MatchCase, ModModule, Parameter, Parameters, Pattern, Stmt, Suite,
-    UnaryOp,
-};
-use ruff_python_parser::Parsed;
-use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use ruff_diagnostics::{Diagnostic, IsolationLevel};
 use ruff_notebook::{CellOffsets, NotebookIndex};
@@ -47,10 +40,16 @@ use ruff_python_ast::identifier::Identifier;
 use ruff_python_ast::name::QualifiedName;
 use ruff_python_ast::str::Quote;
 use ruff_python_ast::visitor::{walk_except_handler, walk_pattern, Visitor};
+use ruff_python_ast::{
+    self as ast, AnyParameterRef, Comprehension, ElifElseClause, ExceptHandler, Expr, ExprContext,
+    FStringElement, Keyword, MatchCase, ModExpression, ModModule, Parameter, Parameters, Pattern,
+    Stmt, Suite, UnaryOp,
+};
 use ruff_python_ast::{helpers, str, visitor, PySourceType};
 use ruff_python_codegen::{Generator, Stylist};
 use ruff_python_index::Indexer;
 use ruff_python_parser::typing::{parse_type_annotation, AnnotationKind};
+use ruff_python_parser::Parsed;
 use ruff_python_semantic::all::{DunderAllDefinition, DunderAllFlags};
 use ruff_python_semantic::analyze::{imports, typing};
 use ruff_python_semantic::{
@@ -60,6 +59,7 @@ use ruff_python_semantic::{
 };
 use ruff_python_stdlib::builtins::{IPYTHON_BUILTINS, MAGIC_GLOBALS, PYTHON_BUILTINS};
 use ruff_source_file::{Locator, OneIndexed, SourceRow};
+use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::checkers::ast::annotation::AnnotationContext;
 use crate::docstrings::extraction::ExtractionTarget;
@@ -2148,7 +2148,10 @@ impl<'a> Checker<'a> {
     ///
     /// class Bar: pass
     /// ```
-    fn visit_deferred_string_type_definitions(&mut self, allocator: &'a typed_arena::Arena<Expr>) {
+    fn visit_deferred_string_type_definitions(
+        &mut self,
+        allocator: &'a typed_arena::Arena<Parsed<ModExpression>>,
+    ) {
         let snapshot = self.semantic.snapshot();
         while !self.visit.string_type_definitions.is_empty() {
             let type_definitions = std::mem::take(&mut self.visit.string_type_definitions);
@@ -2183,7 +2186,7 @@ impl<'a> Checker<'a> {
 
                     self.semantic.flags |=
                         SemanticModelFlags::TYPE_DEFINITION | type_definition_flag;
-                    self.visit_expr(parsed_annotation);
+                    self.visit_expr(parsed_annotation.expr());
                 } else {
                     if self.enabled(Rule::ForwardAnnotationSyntaxError) {
                         self.diagnostics.push(Diagnostic::new(
@@ -2258,7 +2261,7 @@ impl<'a> Checker<'a> {
     /// After initial traversal of the source tree has been completed,
     /// recursively visit all AST nodes that were deferred on the first pass.
     /// This includes lambdas, functions, type parameters, and type annotations.
-    fn visit_deferred(&mut self, allocator: &'a typed_arena::Arena<Expr>) {
+    fn visit_deferred(&mut self, allocator: &'a typed_arena::Arena<Parsed<ModExpression>>) {
         while !self.visit.is_empty() {
             self.visit_deferred_class_bases();
             self.visit_deferred_functions();

diff --git a/crates/ruff_linter/src/rules/flake8_annotations/rules/definition.rs b/crates/ruff_linter/src/rules/flake8_annotations/rules/definition.rs
@@ -518,7 +518,7 @@ fn check_dynamically_typed<F>(
             parse_type_annotation(string_expr, checker.locator().contents())
         {
             if type_hint_resolves_to_any(
-                &parsed_annotation,
+                parsed_annotation.expr(),
                 checker.semantic(),
                 checker.locator(),
                 checker.settings.target_version.minor(),

diff --git a/crates/ruff_linter/src/rules/ruff/rules/implicit_optional.rs b/crates/ruff_linter/src/rules/ruff/rules/implicit_optional.rs
@@ -183,7 +183,7 @@ pub(crate) fn implicit_optional(checker: &mut Checker, parameters: &Parameters)
                 parse_type_annotation(string_expr, checker.locator().contents())
             {
                 let Some(expr) = type_hint_explicitly_allows_none(
-                    &parsed_annotation,
+                    parsed_annotation.expr(),
                     checker.semantic(),
                     checker.locator(),
                     checker.settings.target_version.minor(),

diff --git a/crates/ruff_linter/src/rules/ruff/typing.rs b/crates/ruff_linter/src/rules/ruff/typing.rs
@@ -112,10 +112,15 @@ impl<'a> TypingTarget<'a> {
                 ..
             }) => Some(TypingTarget::PEP604Union(left, right)),
             Expr::NoneLiteral(_) => Some(TypingTarget::None),
-            Expr::StringLiteral(string_expr) => {
-                parse_type_annotation(string_expr, locator.contents())
-                    .map_or(None, |(expr, _)| Some(TypingTarget::ForwardReference(expr)))
-            }
+            Expr::StringLiteral(string_expr) => parse_type_annotation(
+                string_expr,
+                locator.contents(),
+            )
+            .map_or(None, |(parsed_annotation, _)| {
+                Some(TypingTarget::ForwardReference(
+                    parsed_annotation.into_expr(),
+                ))
+            }),
             _ => semantic.resolve_qualified_name(expr).map_or(
                 // If we can't resolve the call path, it must be defined in the
                 // same file, so we assume it's `Any` as it could be a type alias.

diff --git a/crates/ruff_python_parser/src/lib.rs b/crates/ruff_python_parser/src/lib.rs
@@ -357,6 +357,11 @@ impl Parsed<ModExpression> {
         &self.syntax.body
     }
 
+    /// Returns a mutable reference to the expression contained in this parsed output.
+    pub fn expr_mut(&mut self) -> &mut Expr {
+        &mut self.syntax.body
+    }
+
     /// Consumes the [`Parsed`] output and returns the contained [`Expr`].
     pub fn into_expr(self) -> Expr {
         *self.syntax.body

diff --git a/crates/ruff_python_parser/src/typing.rs b/crates/ruff_python_parser/src/typing.rs
@@ -2,10 +2,10 @@
 
 use ruff_python_ast::relocate::relocate_expr;
 use ruff_python_ast::str::raw_contents;
-use ruff_python_ast::{Expr, ExprStringLiteral, StringFlags, StringLiteral};
+use ruff_python_ast::{ExprStringLiteral, ModExpression, StringFlags, StringLiteral};
 use ruff_text_size::Ranged;
 
-use crate::{parse_expression, parse_expression_range, ParseError};
+use crate::{parse_expression, parse_expression_range, ParseError, Parsed};
 
 #[derive(Copy, Clone, Debug)]
 pub enum AnnotationKind {
@@ -34,7 +34,7 @@ impl AnnotationKind {
 pub fn parse_type_annotation(
     string_expr: &ExprStringLiteral,
     source: &str,
-) -> Result<(Expr, AnnotationKind), ParseError> {
+) -> Result<(Parsed<ModExpression>, AnnotationKind), ParseError> {
     let expr_text = &source[string_expr.range()];
 
     if let [string_literal] = string_expr.value.as_slice() {
@@ -58,24 +58,23 @@ pub fn parse_type_annotation(
 fn parse_simple_type_annotation(
     string_literal: &StringLiteral,
     source: &str,
-) -> Result<(Expr, AnnotationKind), ParseError> {
+) -> Result<(Parsed<ModExpression>, AnnotationKind), ParseError> {
     Ok((
         parse_expression_range(
             source,
             string_literal
                 .range()
                 .add_start(string_literal.flags.opener_len())
                 .sub_end(string_literal.flags.closer_len()),
-        )?
-        .into_expr(),
+        )?,
         AnnotationKind::Simple,
     ))
 }
 
 fn parse_complex_type_annotation(
     string_expr: &ExprStringLiteral,
-) -> Result<(Expr, AnnotationKind), ParseError> {
-    let mut parsed = parse_expression(string_expr.value.to_str())?.into_expr();
-    relocate_expr(&mut parsed, string_expr.range());
+) -> Result<(Parsed<ModExpression>, AnnotationKind), ParseError> {
+    let mut parsed = parse_expression(string_expr.value.to_str())?;
+    relocate_expr(parsed.expr_mut(), string_expr.range());
     Ok((parsed, AnnotationKind::Complex))
 }