[`pycodestyle`] Avoid blank line rules for the first logical line in cell (#10291)

## Summary

Closes #10228

The PR makes the blank lines rules keep track of the cell status when
running on a notebook, and makes the rules not trigger when the line is
the first of the cell.

## Test Plan

The example given in #10228 is added as a fixture, along with a few
tests from the main blank lines fixtures.

crates/ruff_linter/resources/test/fixtures/pycodestyle/E30.ipynb

@@ -0,0 +1,228 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "palRUQyD-U6u"
+   },
+   "outputs": [],
+   "source": [
+    "some_string = \"123123\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "UWdDLRyf-Zz0"
+   },
+   "outputs": [],
+   "source": [
+    "some_computation = 1 + 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "YreT1sTr-c32"
+   },
+   "outputs": [],
+   "source": [
+    "some_computation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "V48ppml7-h0f"
+   },
+   "outputs": [],
+   "source": [
+    "def fn():\n",
+    "    print(\"Hey!\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "cscw_8Xv-lYQ"
+   },
+   "outputs": [],
+   "source": [
+    "fn()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# E301\n",
+    "class Class:\n",
+    "    \"\"\"Class for minimal repo.\"\"\"\n",
+    "\n",
+    "    def method(cls) -> None:\n",
+    "        pass\n",
+    "    @classmethod\n",
+    "    def cls_method(cls) -> None:\n",
+    "        pass\n",
+    "# end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# E302\n",
+    "def a():\n",
+    "    pass\n",
+    "\n",
+    "def b():\n",
+    "    pass\n",
+    "# end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# E303\n",
+    "def fn():\n",
+    "    _ = None\n",
+    "\n",
+    "\n",
+    "    # arbitrary comment\n",
+    "\n",
+    "    def inner():  # E306 not expected (pycodestyle detects E306)\n",
+    "        pass\n",
+    "# end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# E303"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "def fn():\n",
+    "\tpass\n",
+    "# end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# E304\n",
+    "@decorator\n",
+    "\n",
+    "def function():\n",
+    "    pass\n",
+    "# end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# E305:7:1\n",
+    "def fn():\n",
+    "    print()\n",
+    "\n",
+    "    # comment\n",
+    "\n",
+    "    # another comment\n",
+    "fn()\n",
+    "# end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# E306:3:5\n",
+    "def a():\n",
+    "    x = 1\n",
+    "    def b():\n",
+    "        pass\n",
+    "# end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Ok\n",
+    "def function1():\n",
+    "\tpass\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def function2():\n",
+    "\tpass\n",
+    "# end"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "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.11.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

crates/ruff_linter/src/checkers/tokens.rs

@@ -41,7 +41,7 @@ pub(crate) fn check_tokens(
         Rule::BlankLinesAfterFunctionOrClass,
         Rule::BlankLinesBeforeNestedDefinition,
     ]) {
-        BlankLinesChecker::new(locator, stylist, settings, source_type)
+        BlankLinesChecker::new(locator, stylist, settings, source_type, cell_offsets)
             .check_lines(tokens, &mut diagnostics);
     }
 

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

@@ -279,6 +279,22 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(Rule::BlankLineBetweenMethods)]
+    #[test_case(Rule::BlankLinesTopLevel)]
+    #[test_case(Rule::TooManyBlankLines)]
+    #[test_case(Rule::BlankLineAfterDecorator)]
+    #[test_case(Rule::BlankLinesAfterFunctionOrClass)]
+    #[test_case(Rule::BlankLinesBeforeNestedDefinition)]
+    fn blank_lines_notebook(rule_code: Rule) -> Result<()> {
+        let snapshot = format!("blank_lines_{}_notebook", rule_code.noqa_code());
+        let diagnostics = test_path(
+            Path::new("pycodestyle").join("E30.ipynb"),
+            &settings::LinterSettings::for_rule(rule_code),
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
+
     #[test]
     fn blank_lines_typing_stub_isort() -> Result<()> {
         let diagnostics = test_path(

crates/ruff_linter/src/rules/pycodestyle/rules/blank_lines.rs

@@ -1,5 +1,7 @@
 use itertools::Itertools;
+use ruff_notebook::CellOffsets;
 use std::cmp::Ordering;
+use std::iter::Peekable;
 use std::num::NonZeroU32;
 use std::slice::Iter;
 
@@ -359,6 +361,9 @@ struct LogicalLineInfo {
     // `true` if this is not a blank but only consists of a comment.
     is_comment_only: bool,
 
+    /// If running on a notebook, whether the line is the first logical line (or a comment preceding it) of its cell.
+    is_beginning_of_cell: bool,
+
     /// `true` if the line is a string only (including trivia tokens) line, which is a docstring if coming right after a class/function definition.
     is_docstring: bool,
 
@@ -387,6 +392,10 @@ struct LinePreprocessor<'a> {
     /// Maximum number of consecutive blank lines between the current line and the previous non-comment logical line.
     /// One of its main uses is to allow a comment to directly precede a class/function definition.
     max_preceding_blank_lines: BlankLines,
+    /// The cell offsets of the notebook (if running on a notebook).
+    cell_offsets: Option<Peekable<Iter<'a, TextSize>>>,
+    /// If running on a notebook, whether the line is the first logical line (or a comment preceding it) of its cell.
+    is_beginning_of_cell: bool,
 }
 
 impl<'a> LinePreprocessor<'a> {
@@ -394,6 +403,7 @@ impl<'a> LinePreprocessor<'a> {
         tokens: &'a [LexResult],
         locator: &'a Locator,
         indent_width: IndentWidth,
+        cell_offsets: Option<&'a CellOffsets>,
     ) -> LinePreprocessor<'a> {
         LinePreprocessor {
             tokens: tokens.iter(),
@@ -401,6 +411,9 @@ impl<'a> LinePreprocessor<'a> {
             line_start: TextSize::new(0),
             max_preceding_blank_lines: BlankLines::Zero,
             indent_width,
+            is_beginning_of_cell: cell_offsets.is_some(),
+            cell_offsets: cell_offsets
+                .map(|cell_offsets| cell_offsets.get(1..).unwrap_or_default().iter().peekable()),
         }
     }
 }
@@ -435,6 +448,19 @@ impl<'a> Iterator for LinePreprocessor<'a> {
             }
             // At the start of the line...
             else {
+                // Check if we are at the beginning of a cell in a notebook.
+                if let Some(ref mut cell_offsets) = self.cell_offsets {
+                    if cell_offsets
+                        .peek()
+                        .is_some_and(|offset| offset == &&self.line_start)
+                    {
+                        self.is_beginning_of_cell = true;
+                        cell_offsets.next();
+                        blank_lines = BlankLines::Zero;
+                        self.max_preceding_blank_lines = BlankLines::Zero;
+                    }
+                }
+
                 // An empty line
                 if token_kind == TokenKind::NonLogicalNewline {
                     blank_lines.add(*range);
@@ -499,6 +525,7 @@ impl<'a> Iterator for LinePreprocessor<'a> {
                         last_token,
                         logical_line_end: range.end(),
                         is_comment_only: line_is_comment_only,
+                        is_beginning_of_cell: self.is_beginning_of_cell,
                         is_docstring,
                         indent_length,
                         blank_lines,
@@ -513,6 +540,10 @@ impl<'a> Iterator for LinePreprocessor<'a> {
                     // Set the start for the next logical line.
                     self.line_start = range.end();
 
+                    if self.cell_offsets.is_some() && !line_is_comment_only {
+                        self.is_beginning_of_cell = false;
+                    }
+
                     return Some(logical_line);
                 }
                 _ => {}
@@ -662,6 +693,7 @@ pub(crate) struct BlankLinesChecker<'a> {
     lines_after_imports: isize,
     lines_between_types: usize,
     source_type: PySourceType,
+    cell_offsets: Option<&'a CellOffsets>,
 }
 
 impl<'a> BlankLinesChecker<'a> {
@@ -670,6 +702,7 @@ impl<'a> BlankLinesChecker<'a> {
         stylist: &'a Stylist<'a>,
         settings: &crate::settings::LinterSettings,
         source_type: PySourceType,
+        cell_offsets: Option<&'a CellOffsets>,
     ) -> BlankLinesChecker<'a> {
         BlankLinesChecker {
             stylist,
@@ -678,6 +711,7 @@ impl<'a> BlankLinesChecker<'a> {
             lines_after_imports: settings.isort.lines_after_imports,
             lines_between_types: settings.isort.lines_between_types,
             source_type,
+            cell_offsets,
         }
     }
 
@@ -685,7 +719,8 @@ impl<'a> BlankLinesChecker<'a> {
     pub(crate) fn check_lines(&self, tokens: &[LexResult], diagnostics: &mut Vec<Diagnostic>) {
         let mut prev_indent_length: Option<usize> = None;
         let mut state = BlankLinesState::default();
-        let line_preprocessor = LinePreprocessor::new(tokens, self.locator, self.indent_width);
+        let line_preprocessor =
+            LinePreprocessor::new(tokens, self.locator, self.indent_width, self.cell_offsets);
 
         for logical_line in line_preprocessor {
             // Reset `follows` after a dedent:
@@ -826,6 +861,8 @@ impl<'a> BlankLinesChecker<'a> {
             && !self.source_type.is_stub()
             // Do not expect blank lines before the first logical line.
             && state.is_not_first_logical_line
+            // Ignore the first logical line (and any comment preceding it) of each cell in notebooks.
+            && !line.is_beginning_of_cell
         {
             // E302
             let mut diagnostic = Diagnostic::new(
@@ -940,6 +977,8 @@ impl<'a> BlankLinesChecker<'a> {
             && !line.kind.is_class_function_or_decorator()
             // Blank lines in stub files are used for grouping, don't enforce blank lines.
             && !self.source_type.is_stub()
+            // Ignore the first logical line (and any comment preceding it) of each cell in notebooks.
+            && !line.is_beginning_of_cell
         {
             // E305
             let mut diagnostic = Diagnostic::new(

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__blank_lines_E301_notebook.snap

@@ -0,0 +1,22 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+---
+E30.ipynb:13:5: E301 [*] Expected 1 blank line, found 0
+   |
+11 |     def method(cls) -> None:
+12 |         pass
+13 |     @classmethod
+   |     ^ E301
+14 |     def cls_method(cls) -> None:
+15 |         pass
+   |
+   = help: Add missing blank line
+
+ℹ Safe fix
+10 10 | 
+11 11 |     def method(cls) -> None:
+12 12 |         pass
+   13 |+
+13 14 |     @classmethod
+14 15 |     def cls_method(cls) -> None:
+15 16 |         pass

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__blank_lines_E302_notebook.snap

@@ -0,0 +1,22 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+---
+E30.ipynb:21:1: E302 [*] Expected 2 blank lines, found 1
+   |
+19 |     pass
+20 | 
+21 | def b():
+   | ^^^ E302
+22 |     pass
+23 | # end
+   |
+   = help: Add missing blank line(s)
+
+ℹ Safe fix
+18 18 | def a():
+19 19 |     pass
+20 20 | 
+   21 |+
+21 22 | def b():
+22 23 |     pass
+23 24 | # end

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__blank_lines_E303_notebook.snap

@@ -0,0 +1,39 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+---
+E30.ipynb:29:5: E303 [*] Too many blank lines (2)
+   |
+29 |     # arbitrary comment
+   |     ^^^^^^^^^^^^^^^^^^^ E303
+30 | 
+31 |     def inner():  # E306 not expected (pycodestyle detects E306)
+   |
+   = help: Remove extraneous blank line(s)
+
+ℹ Safe fix
+25 25 | def fn():
+26 26 |     _ = None
+27 27 | 
+28    |-
+29 28 |     # arbitrary comment
+30 29 | 
+31 30 |     def inner():  # E306 not expected (pycodestyle detects E306)
+
+E30.ipynb:39:1: E303 [*] Too many blank lines (4)
+   |
+39 | def fn():
+   | ^^^ E303
+40 |     pass
+41 | # end
+   |
+   = help: Remove extraneous blank line(s)
+
+ℹ Safe fix
+34 34 | # E303
+35 35 | 
+36 36 | 
+37    |-
+38    |-
+39 37 | def fn():
+40 38 | 	pass
+41 39 | # end

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__blank_lines_E304_notebook.snap

@@ -0,0 +1,22 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+---
+E30.ipynb:45:1: E304 [*] Blank lines found after function decorator (1)
+   |
+43 | @decorator
+44 | 
+45 | def function():
+   | ^^^ E304
+46 |     pass
+47 | # end
+   |
+   = help: Remove extraneous blank line(s)
+
+ℹ Safe fix
+41 41 | # end
+42 42 | # E304
+43 43 | @decorator
+44    |-
+45 44 | def function():
+46 45 |     pass
+47 46 | # end

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__blank_lines_E305_notebook.snap

@@ -0,0 +1,22 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+---
+E30.ipynb:55:1: E305 [*] Expected 2 blank lines after class or function definition, found (1)
+   |
+54 |     # another comment
+55 | fn()
+   | ^^ E305
+56 | # end
+57 | # E306:3:5
+   |
+   = help: Add missing blank line(s)
+
+ℹ Safe fix
+52 52 |     # comment
+53 53 | 
+54 54 |     # another comment
+   55 |+
+   56 |+
+55 57 | fn()
+56 58 | # end
+57 59 | # E306:3:5

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__blank_lines_E306_notebook.snap

@@ -0,0 +1,22 @@
+---
+source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
+---
+E30.ipynb:60:5: E306 [*] Expected 1 blank line before a nested definition, found 0
+   |
+58 | def a():
+59 |     x = 1
+60 |     def b():
+   |     ^^^ E306
+61 |         pass
+62 | # end
+   |
+   = help: Add missing blank line
+
+ℹ Safe fix
+57 57 | # E306:3:5
+58 58 | def a():
+59 59 |     x = 1
+   60 |+
+60 61 |     def b():
+61 62 |         pass
+62 63 | # end