[syntax-errors] Async comprehension in sync comprehension (#17177)

Summary
--

Detect async comprehensions nested in sync comprehensions in async
functions before Python 3.11, when this was [changed].

The actual logic of this rule is very straightforward, but properly
tracking the async scopes took a bit of work. An alternative to the
current approach is to offload the `in_async_context` check into the
`SemanticSyntaxContext` trait, but that actually required much more
extensive changes to the `TestContext` and also to ruff's semantic
model, as you can see in the changes up to
31554b4. This version has the benefit
of mostly centralizing the state tracking in `SemanticSyntaxChecker`,
although there was some subtlety around deferred function body traversal
that made the changes to `Checker` more intrusive too (hence the new
linter test).

The `Checkpoint` struct/system is obviously overkill for now since it's
only tracking a single `bool`, but I thought it might be more useful
later.

[changed]: python/cpython#77527

Test Plan
--

New inline tests and a new linter integration test.

Author: ntBre

crates/ruff_linter/resources/test/fixtures/syntax_errors/async_comprehension.ipynb

@@ -0,0 +1,57 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a70904e-dbfe-441c-99ec-12e6cf57f8ba",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "async def elements(n): yield n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5412fc2f-76eb-42c0-8db1-b5af6fdc46aa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "[x async for x in elements(5)]  # okay, async at top level"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dc3c94a7-2e64-42de-9351-260b3f41c3fd",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "[[x async for x in elements(5)] for i in range(5)]  # error on 3.10, okay after"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -27,7 +27,8 @@ use std::path::Path;
 use itertools::Itertools;
 use log::debug;
 use ruff_python_parser::semantic_errors::{
-    SemanticSyntaxChecker, SemanticSyntaxContext, SemanticSyntaxError, SemanticSyntaxErrorKind,
+    Checkpoint, SemanticSyntaxChecker, SemanticSyntaxContext, SemanticSyntaxError,
+    SemanticSyntaxErrorKind,
 };
 use rustc_hash::{FxHashMap, FxHashSet};
 
@@ -282,7 +283,7 @@ impl<'a> Checker<'a> {
             last_stmt_end: TextSize::default(),
             docstring_state: DocstringState::default(),
             target_version,
-            semantic_checker: SemanticSyntaxChecker::new(),
+            semantic_checker: SemanticSyntaxChecker::new(source_type),
             semantic_errors: RefCell::default(),
         }
     }
@@ -525,10 +526,14 @@ impl<'a> Checker<'a> {
         self.target_version
     }
 
-    fn with_semantic_checker(&mut self, f: impl FnOnce(&mut SemanticSyntaxChecker, &Checker)) {
+    fn with_semantic_checker(
+        &mut self,
+        f: impl FnOnce(&mut SemanticSyntaxChecker, &Checker) -> Checkpoint,
+    ) -> Checkpoint {
         let mut checker = std::mem::take(&mut self.semantic_checker);
-        f(&mut checker, self);
+        let checkpoint = f(&mut checker, self);
         self.semantic_checker = checker;
+        checkpoint
     }
 }
 
@@ -576,7 +581,8 @@ impl SemanticSyntaxContext for Checker<'_> {
             | SemanticSyntaxErrorKind::InvalidExpression(..)
             | SemanticSyntaxErrorKind::DuplicateMatchKey(_)
             | SemanticSyntaxErrorKind::DuplicateMatchClassAttribute(_)
-            | SemanticSyntaxErrorKind::InvalidStarExpression => {
+            | SemanticSyntaxErrorKind::InvalidStarExpression
+            | SemanticSyntaxErrorKind::AsyncComprehensionOutsideAsyncFunction(_) => {
                 if self.settings.preview.is_enabled() {
                     self.semantic_errors.borrow_mut().push(error);
                 }
@@ -595,7 +601,13 @@ impl SemanticSyntaxContext for Checker<'_> {
 
 impl<'a> Visitor<'a> for Checker<'a> {
     fn visit_stmt(&mut self, stmt: &'a Stmt) {
-        self.with_semantic_checker(|semantic, context| semantic.visit_stmt(stmt, context));
+        // For functions, defer semantic syntax error checks until the body of the function is
+        // visited
+        let checkpoint = if stmt.is_function_def_stmt() {
+            None
+        } else {
+            Some(self.with_semantic_checker(|semantic, context| semantic.enter_stmt(stmt, context)))
+        };
 
         // Step 0: Pre-processing
         self.semantic.push_node(stmt);
@@ -1198,6 +1210,10 @@ impl<'a> Visitor<'a> for Checker<'a> {
         self.semantic.flags = flags_snapshot;
         self.semantic.pop_node();
         self.last_stmt_end = stmt.end();
+
+        if let Some(checkpoint) = checkpoint {
+            self.semantic_checker.exit_stmt(checkpoint);
+        }
     }
 
     fn visit_annotation(&mut self, expr: &'a Expr) {
@@ -1208,7 +1224,8 @@ impl<'a> Visitor<'a> for Checker<'a> {
     }
 
     fn visit_expr(&mut self, expr: &'a Expr) {
-        self.with_semantic_checker(|semantic, context| semantic.visit_expr(expr, context));
+        let checkpoint =
+            self.with_semantic_checker(|semantic, context| semantic.enter_expr(expr, context));
 
         // Step 0: Pre-processing
         if self.source_type.is_stub()
@@ -1755,6 +1772,8 @@ impl<'a> Visitor<'a> for Checker<'a> {
         self.semantic.flags = flags_snapshot;
         analyze::expression(expr, self);
         self.semantic.pop_node();
+
+        self.semantic_checker.exit_expr(checkpoint);
     }
 
     fn visit_except_handler(&mut self, except_handler: &'a ExceptHandler) {
@@ -2590,17 +2609,24 @@ impl<'a> Checker<'a> {
             for snapshot in deferred_functions {
                 self.semantic.restore(snapshot);
 
+                let stmt = self.semantic.current_statement();
+
                 let Stmt::FunctionDef(ast::StmtFunctionDef {
                     body, parameters, ..
-                }) = self.semantic.current_statement()
+                }) = stmt
                 else {
                     unreachable!("Expected Stmt::FunctionDef")
                 };
 
+                let checkpoint = self
+                    .with_semantic_checker(|semantic, context| semantic.enter_stmt(stmt, context));
+
                 self.visit_parameters(parameters);
                 // Set the docstring state before visiting the function body.
                 self.docstring_state = DocstringState::Expected(ExpectedDocstringKind::Function);
                 self.visit_body(body);
+
+                self.semantic_checker.exit_stmt(checkpoint);
             }
         }
         self.semantic.restore(snapshot);

crates/ruff_linter/src/linter.rs

@@ -777,14 +777,22 @@ mod tests {
     use std::path::Path;
 
     use anyhow::Result;
+    use ruff_python_ast::{PySourceType, PythonVersion};
+    use ruff_python_codegen::Stylist;
+    use ruff_python_index::Indexer;
+    use ruff_python_parser::ParseOptions;
+    use ruff_python_trivia::textwrap::dedent;
+    use ruff_text_size::Ranged;
     use test_case::test_case;
 
     use ruff_notebook::{Notebook, NotebookError};
 
+    use crate::linter::check_path;
+    use crate::message::Message;
     use crate::registry::Rule;
     use crate::source_kind::SourceKind;
     use crate::test::{assert_notebook_path, test_contents, TestedNotebook};
-    use crate::{assert_messages, settings};
+    use crate::{assert_messages, directives, settings, Locator};
 
     /// Construct a path to a Jupyter notebook in the `resources/test/fixtures/jupyter` directory.
     fn notebook_path(path: impl AsRef<Path>) -> std::path::PathBuf {
@@ -934,4 +942,122 @@ mod tests {
         }
         Ok(())
     }
+
+    /// Wrapper around `test_contents_syntax_errors` for testing a snippet of code instead of a
+    /// file.
+    fn test_snippet_syntax_errors(
+        contents: &str,
+        settings: &settings::LinterSettings,
+    ) -> Vec<Message> {
+        let contents = dedent(contents);
+        test_contents_syntax_errors(
+            &SourceKind::Python(contents.to_string()),
+            Path::new("<filename>"),
+            settings,
+        )
+    }
+
+    /// A custom test runner that prints syntax errors in addition to other diagnostics. Adapted
+    /// from `flakes` in pyflakes/mod.rs.
+    fn test_contents_syntax_errors(
+        source_kind: &SourceKind,
+        path: &Path,
+        settings: &settings::LinterSettings,
+    ) -> Vec<Message> {
+        let source_type = PySourceType::from(path);
+        let options =
+            ParseOptions::from(source_type).with_target_version(settings.unresolved_target_version);
+        let parsed = ruff_python_parser::parse_unchecked(source_kind.source_code(), options)
+            .try_into_module()
+            .expect("PySourceType always parses into a module");
+        let locator = Locator::new(source_kind.source_code());
+        let stylist = Stylist::from_tokens(parsed.tokens(), locator.contents());
+        let indexer = Indexer::from_tokens(parsed.tokens(), locator.contents());
+        let directives = directives::extract_directives(
+            parsed.tokens(),
+            directives::Flags::from_settings(settings),
+            &locator,
+            &indexer,
+        );
+        let mut messages = check_path(
+            path,
+            None,
+            &locator,
+            &stylist,
+            &indexer,
+            &directives,
+            settings,
+            settings::flags::Noqa::Enabled,
+            source_kind,
+            source_type,
+            &parsed,
+            settings.unresolved_target_version,
+        );
+        messages.sort_by_key(Ranged::start);
+        messages
+    }
+
+    #[test_case(
+        "error_on_310",
+        "async def f(): return [[x async for x in foo(n)] for n in range(3)]",
+        PythonVersion::PY310
+    )]
+    #[test_case(
+        "okay_on_311",
+        "async def f(): return [[x async for x in foo(n)] for n in range(3)]",
+        PythonVersion::PY311
+    )]
+    #[test_case(
+        "okay_on_310",
+        "async def test(): return [[x async for x in elements(n)] async for n in range(3)]",
+        PythonVersion::PY310
+    )]
+    #[test_case(
+        "deferred_function_body",
+        "
+		async def f(): [x for x in foo()] and [x async for x in foo()]
+		async def f():
+			def g(): ...
+			[x async for x in foo()]
+		",
+        PythonVersion::PY310
+    )]
+    fn test_async_comprehension_in_sync_comprehension(
+        name: &str,
+        contents: &str,
+        python_version: PythonVersion,
+    ) {
+        let snapshot = format!("async_comprehension_in_sync_comprehension_{name}_{python_version}");
+        let messages = test_snippet_syntax_errors(
+            contents,
+            &settings::LinterSettings {
+                rules: settings::rule_table::RuleTable::empty(),
+                unresolved_target_version: python_version,
+                preview: settings::types::PreviewMode::Enabled,
+                ..Default::default()
+            },
+        );
+        assert_messages!(snapshot, messages);
+    }
+
+    #[test_case(PythonVersion::PY310)]
+    #[test_case(PythonVersion::PY311)]
+    fn test_async_comprehension_notebook(python_version: PythonVersion) -> Result<()> {
+        let snapshot =
+            format!("async_comprehension_in_sync_comprehension_notebook_{python_version}");
+        let path = Path::new("resources/test/fixtures/syntax_errors/async_comprehension.ipynb");
+        let messages = test_contents_syntax_errors(
+            &SourceKind::IpyNotebook(Notebook::from_path(path)?),
+            path,
+            &settings::LinterSettings {
+                unresolved_target_version: python_version,
+                rules: settings::rule_table::RuleTable::empty(),
+                preview: settings::types::PreviewMode::Enabled,
+                ..Default::default()
+            },
+        );
+        assert_messages!(snapshot, messages);
+
+        Ok(())
+    }
 }

crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__async_comprehension_in_sync_comprehension_deferred_function_body_3.10.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/linter.rs
+---
+

crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__async_comprehension_in_sync_comprehension_error_on_310_3.10.snap

@@ -0,0 +1,8 @@
+---
+source: crates/ruff_linter/src/linter.rs
+---
+<filename>:1:27: SyntaxError: cannot use an asynchronous comprehension outside of an asynchronous function on Python 3.10 (syntax was added in 3.11)
+  |
+1 | async def f(): return [[x async for x in foo(n)] for n in range(3)]
+  |                           ^^^^^^^^^^^^^^^^^^^^^
+  |

crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__async_comprehension_in_sync_comprehension_notebook_3.10.snap

@@ -0,0 +1,10 @@
+---
+source: crates/ruff_linter/src/linter.rs
+---
+resources/test/fixtures/syntax_errors/async_comprehension.ipynb:3:5: SyntaxError: cannot use an asynchronous comprehension outside of an asynchronous function on Python 3.10 (syntax was added in 3.11)
+  |
+1 | async def elements(n): yield n
+2 | [x async for x in elements(5)]  # okay, async at top level
+3 | [[x async for x in elements(5)] for i in range(5)]  # error on 3.10, okay after
+  |     ^^^^^^^^^^^^^^^^^^^^^^^^^^
+  |

crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__async_comprehension_in_sync_comprehension_notebook_3.11.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/linter.rs
+---
+

crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__async_comprehension_in_sync_comprehension_okay_on_310_3.10.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/linter.rs
+---
+

crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__async_comprehension_in_sync_comprehension_okay_on_311_3.11.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/linter.rs
+---
+

crates/ruff_python_parser/resources/inline/err/nested_async_comprehension_py310.py

@@ -0,0 +1,6 @@
+# parse_options: {"target-version": "3.10"}
+async def f(): return [[x async for x in foo(n)] for n in range(3)]    # list
+async def g(): return [{x: 1 async for x in foo(n)} for n in range(3)] # dict
+async def h(): return [{x async for x in foo(n)} for n in range(3)]    # set
+async def i(): return [([y async for y in range(1)], [z for z in range(2)]) for x in range(5)]
+async def j(): return [([y for y in range(1)], [z async for z in range(2)]) for x in range(5)]