[flake8-pyi] Avoid syntax error from conflict with PIE790 (PYI021) (#20010)

<!--
Thank you for contributing to Ruff/ty! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title? (Please prefix
with `[ty]` for ty pull
  requests.)
- Does this pull request include references to any relevant issues?
-->

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->
First contribution so please let me know if I've made a mistake
anywhere. This was aimed to fix #19982, it adds the isolation level to
PYI021 to in the same style as the PIE790 rule.

fixes: #19982

## Test Plan

<!-- How was it tested? -->
I added a case to the PYI021.pyi file where the two rules are present as
there wasn't a case with them both interacting, using the minimal
reproducible example that @ntBre created on the issue (I think I got the
`# ERROR` markings wrong, so please let me know how to fix that if I
did).

---------

Co-authored-by: Brent Westbrook <brentrwestbrook@gmail.com>

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI021_1.pyi

@@ -0,0 +1,18 @@
+from contextlib import nullcontext
+
+
+def check_isolation_level(mode: int) -> None:
+    """Will report both, but only fix the first."""  # ERROR PYI021
+    ...  # ERROR PIE790
+
+
+with nullcontext():
+    """Should not report."""
+    # add something thats not a pass here
+    # to not get a remove unnecessary pass err
+    x = 0
+
+if True:
+    """Should not report."""
+    # same as above
+    y = 1

crates/ruff_linter/src/checkers/ast/analyze/definitions.rs

@@ -145,9 +145,6 @@ pub(crate) fn definitions(checker: &mut Checker) {
         }
 
         // flake8-pyi
-        if enforce_stubs {
-            flake8_pyi::rules::docstring_in_stubs(checker, definition, docstring);
-        }
         if enforce_stubs_and_runtime {
             flake8_pyi::rules::iter_method_return_iterable(checker, definition);
         }

crates/ruff_linter/src/checkers/ast/analyze/suite.rs

@@ -2,14 +2,17 @@ use ruff_python_ast::Stmt;
 
 use crate::checkers::ast::Checker;
 use crate::codes::Rule;
-use crate::rules::flake8_pie;
 use crate::rules::refurb;
+use crate::rules::{flake8_pie, flake8_pyi};
 
 /// Run lint rules over a suite of [`Stmt`] syntax nodes.
 pub(crate) fn suite(suite: &[Stmt], checker: &Checker) {
     if checker.is_rule_enabled(Rule::UnnecessaryPlaceholder) {
         flake8_pie::rules::unnecessary_placeholder(checker, suite);
     }
+    if checker.source_type.is_stub() && checker.is_rule_enabled(Rule::DocstringInStub) {
+        flake8_pyi::rules::docstring_in_stubs(checker, suite);
+    }
     if checker.is_rule_enabled(Rule::RepeatedGlobal) {
         refurb::rules::repeated_global(checker, suite);
     }

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

@@ -87,7 +87,7 @@ mod deferred;
 
 /// State representing whether a docstring is expected or not for the next statement.
 #[derive(Debug, Copy, Clone, PartialEq)]
-enum DocstringState {
+pub(crate) enum DocstringState {
     /// The next statement is expected to be a docstring, but not necessarily so.
     ///
     /// For example, in the following code:
@@ -128,7 +128,7 @@ impl DocstringState {
 
 /// The kind of an expected docstring.
 #[derive(Debug, Copy, Clone, PartialEq)]
-enum ExpectedDocstringKind {
+pub(crate) enum ExpectedDocstringKind {
     /// A module-level docstring.
     ///
     /// For example,
@@ -603,6 +603,11 @@ impl<'a> Checker<'a> {
     pub(crate) const fn context(&self) -> &'a LintContext<'a> {
         self.context
     }
+
+    /// Return the current [`DocstringState`].
+    pub(crate) fn docstring_state(&self) -> DocstringState {
+        self.docstring_state
+    }
 }
 
 pub(crate) struct TypingImporter<'a, 'b> {

crates/ruff_linter/src/rules/flake8_pyi/mod.rs

@@ -192,4 +192,17 @@ mod tests {
         assert_diagnostics!(snapshot, diagnostics);
         Ok(())
     }
+
+    #[test_case(Path::new("PYI021_1.pyi"))]
+    fn pyi021_pie790_isolation_check(path: &Path) -> Result<()> {
+        let diagnostics = test_path(
+            Path::new("flake8_pyi").join(path).as_path(),
+            &settings::LinterSettings::for_rules([
+                Rule::DocstringInStub,
+                Rule::UnnecessaryPlaceholder,
+            ]),
+        )?;
+        assert_diagnostics!(diagnostics);
+        Ok(())
+    }
 }

crates/ruff_linter/src/rules/flake8_pyi/rules/docstring_in_stubs.rs

@@ -1,9 +1,9 @@
 use ruff_macros::{ViolationMetadata, derive_message_formats};
-use ruff_python_ast::ExprStringLiteral;
-use ruff_python_semantic::Definition;
+use ruff_python_ast::{ExprStringLiteral, Stmt};
 use ruff_text_size::Ranged;
 
-use crate::checkers::ast::Checker;
+use crate::checkers::ast::{Checker, DocstringState, ExpectedDocstringKind};
+use crate::docstrings::extraction::docstring_from;
 use crate::{AlwaysFixableViolation, Edit, Fix};
 
 /// ## What it does
@@ -41,26 +41,34 @@ impl AlwaysFixableViolation for DocstringInStub {
 }
 
 /// PYI021
-pub(crate) fn docstring_in_stubs(
-    checker: &Checker,
-    definition: &Definition,
-    docstring: Option<&ExprStringLiteral>,
-) {
+pub(crate) fn docstring_in_stubs(checker: &Checker, body: &[Stmt]) {
+    if !matches!(
+        checker.docstring_state(),
+        DocstringState::Expected(
+            ExpectedDocstringKind::Module
+                | ExpectedDocstringKind::Class
+                | ExpectedDocstringKind::Function
+        )
+    ) {
+        return;
+    }
+
+    let docstring = docstring_from(body);
+
     let Some(docstring_range) = docstring.map(ExprStringLiteral::range) else {
         return;
     };
 
-    let statements = match definition {
-        Definition::Module(module) => module.python_ast,
-        Definition::Member(member) => member.body(),
-    };
-
-    let edit = if statements.len() == 1 {
+    let edit = if body.len() == 1 {
         Edit::range_replacement("...".to_string(), docstring_range)
     } else {
         Edit::range_deletion(docstring_range)
     };
 
-    let mut diagnostic = checker.report_diagnostic(DocstringInStub, docstring_range);
-    diagnostic.set_fix(Fix::unsafe_edit(edit));
+    let isolation_level = Checker::isolation(checker.semantic().current_statement_id());
+    let fix = Fix::unsafe_edit(edit).isolate(isolation_level);
+
+    checker
+        .report_diagnostic(DocstringInStub, docstring_range)
+        .set_fix(fix);
 }

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__pyi021_pie790_isolation_check.snap

@@ -0,0 +1,39 @@
+---
+source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
+---
+PYI021 [*] Docstrings should not be included in stubs
+ --> PYI021_1.pyi:5:5
+  |
+4 | def check_isolation_level(mode: int) -> None:
+5 |     """Will report both, but only fix the first."""  # ERROR PYI021
+  |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+6 |     ...  # ERROR PIE790
+  |
+help: Remove docstring
+2 | 
+3 | 
+4 | def check_isolation_level(mode: int) -> None:
+  -     """Will report both, but only fix the first."""  # ERROR PYI021
+5 +       # ERROR PYI021
+6 |     ...  # ERROR PIE790
+7 | 
+8 | 
+note: This is an unsafe fix and may change runtime behavior
+
+PIE790 [*] Unnecessary `...` literal
+ --> PYI021_1.pyi:6:5
+  |
+4 | def check_isolation_level(mode: int) -> None:
+5 |     """Will report both, but only fix the first."""  # ERROR PYI021
+6 |     ...  # ERROR PIE790
+  |     ^^^
+  |
+help: Remove unnecessary `...`
+3 | 
+4 | def check_isolation_level(mode: int) -> None:
+5 |     """Will report both, but only fix the first."""  # ERROR PYI021
+  -     ...  # ERROR PIE790
+6 +     # ERROR PIE790
+7 | 
+8 | 
+9 | with nullcontext():