From 64dee42bc6e1fe0f7f1a21b3dd268de27e16fea3 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 26 Jul 2023 12:49:32 +0530
Subject: [PATCH 01/18] Use Jupyter mode for the parser with Notebook files

---
 .../fixtures/jupyter/cell/cell_magic.json     |  8 ++++
 .../test/fixtures/jupyter/line_magics.ipynb   | 41 ++++++++++++++++++
 .../jupyter/line_magics_expected.ipynb        | 40 +++++++++++++++++
 crates/ruff/src/importer/insertion.rs         | 11 +++--
 crates/ruff/src/jupyter/notebook.rs           | 43 +++++++++++--------
 ...jupyter__notebook__tests__line_magics.snap | 21 +++++++++
 crates/ruff/src/linter.rs                     | 26 ++++++++---
 crates/ruff/src/rules/pyflakes/mod.rs         |  3 +-
 crates/ruff/src/test.rs                       | 10 ++++-
 crates/ruff_dev/src/print_ast.rs              | 12 +++++-
 crates/ruff_dev/src/print_tokens.rs           | 10 ++++-
 crates/ruff_python_parser/src/lib.rs          | 12 ++++--
 crates/ruff_shrinking/src/main.rs             |  9 ++--
 crates/ruff_wasm/src/lib.rs                   |  2 +-
 14 files changed, 205 insertions(+), 43 deletions(-)
 create mode 100644 crates/ruff/resources/test/fixtures/jupyter/cell/cell_magic.json
 create mode 100644 crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb
 create mode 100644 crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb
 create mode 100644 crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap

diff --git a/crates/ruff/resources/test/fixtures/jupyter/cell/cell_magic.json b/crates/ruff/resources/test/fixtures/jupyter/cell/cell_magic.json
new file mode 100644
index 0000000000000..ef68b202e6811
--- /dev/null
+++ b/crates/ruff/resources/test/fixtures/jupyter/cell/cell_magic.json
@@ -0,0 +1,8 @@
+{
+  "execution_count": null,
+  "cell_type": "code",
+  "id": "1",
+  "metadata": {},
+  "outputs": [],
+  "source": ["%%timeit\n", "print('hello world')"]
+}
diff --git a/crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb b/crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb
new file mode 100644
index 0000000000000..785a73102ad73
--- /dev/null
+++ b/crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb
@@ -0,0 +1,41 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eab4754a-d6df-4b41-8ee8-7e23aef440f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "\n",
+    "%matplotlib inline\n",
+    "\n",
+    "import os\n",
+    "\n",
+    "_ = math.pi"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (ruff)",
+   "language": "python",
+   "name": "ruff"
+  },
+  "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.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb b/crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb
new file mode 100644
index 0000000000000..cdf69fa719600
--- /dev/null
+++ b/crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb
@@ -0,0 +1,40 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cad32845-44f9-4a53-8b8c-a6b1bb3f3378",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "\n",
+    "%matplotlib inline\n",
+    "\n",
+    "\n",
+    "_ = math.pi"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (ruff)",
+   "language": "python",
+   "name": "ruff"
+  },
+  "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.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/crates/ruff/src/importer/insertion.rs b/crates/ruff/src/importer/insertion.rs
index 5739955bedba2..176fef854e1bd 100644
--- a/crates/ruff/src/importer/insertion.rs
+++ b/crates/ruff/src/importer/insertion.rs
@@ -300,12 +300,11 @@ fn match_leading_semicolon(s: &str) -> Option<TextSize> {
 mod tests {
     use anyhow::Result;
 
-    use ruff_python_parser::lexer::LexResult;
-    use ruff_text_size::TextSize;
-
     use ruff_python_codegen::Stylist;
-    use ruff_python_parser::parse_suite;
+    use ruff_python_parser::lexer::LexResult;
+    use ruff_python_parser::{parse_suite, Mode};
     use ruff_source_file::{LineEnding, Locator};
+    use ruff_text_size::TextSize;
 
     use super::Insertion;
 
@@ -313,7 +312,7 @@ mod tests {
     fn start_of_file() -> Result<()> {
         fn insert(contents: &str) -> Result<Insertion> {
             let program = parse_suite(contents, "<filename>")?;
-            let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents);
+            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             Ok(Insertion::start_of_file(&program, &locator, &stylist))
@@ -424,7 +423,7 @@ x = 1
     #[test]
     fn start_of_block() {
         fn insert(contents: &str, offset: TextSize) -> Insertion {
-            let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents);
+            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             Insertion::start_of_block(offset, &locator, &stylist)
diff --git a/crates/ruff/src/jupyter/notebook.rs b/crates/ruff/src/jupyter/notebook.rs
index 75a7664d712ca..d234c7f17ff0c 100644
--- a/crates/ruff/src/jupyter/notebook.rs
+++ b/crates/ruff/src/jupyter/notebook.rs
@@ -7,6 +7,7 @@ use std::path::Path;
 
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
+use rustpython_parser::Mode;
 use serde::Serialize;
 use serde_json::error::Category;
 
@@ -24,8 +25,6 @@ use crate::IOError;
 
 pub const JUPYTER_NOTEBOOK_EXT: &str = "ipynb";
 
-const MAGIC_PREFIX: [&str; 3] = ["%", "!", "?"];
-
 /// Run round-trip source code generation on a given Jupyter notebook file path.
 pub fn round_trip(path: &Path) -> anyhow::Result<String> {
     let mut notebook = Notebook::read(path).map_err(|err| {
@@ -78,26 +77,21 @@ impl Cell {
     /// Return `true` if it's a valid code cell.
     ///
     /// A valid code cell is a cell where the cell type is [`Cell::Code`] and the
-    /// source doesn't contain a magic, shell or help command.
+    /// source doesn't contain a cell magic.
     fn is_valid_code_cell(&self) -> bool {
         let source = match self {
             Cell::Code(cell) => &cell.source,
             _ => return false,
         };
-        // Ignore a cell if it contains a magic command. There could be valid
-        // Python code as well, but we'll ignore that for now.
-        // TODO(dhruvmanila): https://github.com/psf/black/blob/main/src/black/handle_ipynb_magics.py
+        // Ignore cells containing cell magic. This is different from line magic
+        // which is allowed and ignored by the parser.
         !match source {
-            SourceValue::String(string) => string.lines().any(|line| {
-                MAGIC_PREFIX
-                    .iter()
-                    .any(|prefix| line.trim_start().starts_with(prefix))
-            }),
-            SourceValue::StringArray(string_array) => string_array.iter().any(|line| {
-                MAGIC_PREFIX
-                    .iter()
-                    .any(|prefix| line.trim_start().starts_with(prefix))
-            }),
+            SourceValue::String(string) => string
+                .lines()
+                .any(|line| line.trim_start().starts_with("%%")),
+            SourceValue::StringArray(string_array) => string_array
+                .iter()
+                .any(|line| line.trim_start().starts_with("%%")),
         }
     }
 }
@@ -513,9 +507,10 @@ mod tests {
     }
 
     #[test_case(Path::new("markdown.json"), false; "markdown")]
-    #[test_case(Path::new("only_magic.json"), false; "only_magic")]
-    #[test_case(Path::new("code_and_magic.json"), false; "code_and_magic")]
+    #[test_case(Path::new("only_magic.json"), true; "only_magic")]
+    #[test_case(Path::new("code_and_magic.json"), true; "code_and_magic")]
     #[test_case(Path::new("only_code.json"), true; "only_code")]
+    #[test_case(Path::new("cell_magic.json"), false; "cell_magic")]
     fn test_is_valid_code_cell(path: &Path, expected: bool) -> Result<()> {
         assert_eq!(read_jupyter_cell(path)?.is_valid_code_cell(), expected);
         Ok(())
@@ -576,6 +571,18 @@ print("after empty cells")
         Ok(())
     }
 
+    #[test]
+    fn test_line_magics() -> Result<()> {
+        let path = "line_magics.ipynb".to_string();
+        let (diagnostics, source_kind) = test_notebook_path(
+            &path,
+            Path::new("line_magics_expected.ipynb"),
+            &settings::Settings::for_rule(Rule::UnusedImport),
+        )?;
+        assert_messages!(diagnostics, path, source_kind);
+        Ok(())
+    }
+
     #[test]
     fn test_json_consistency() -> Result<()> {
         let path = "before_fix.ipynb".to_string();
diff --git a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap
new file mode 100644
index 0000000000000..576b938b6b2e0
--- /dev/null
+++ b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap
@@ -0,0 +1,21 @@
+---
+source: crates/ruff/src/jupyter/notebook.rs
+---
+line_magics.ipynb:cell 1:5:8: F401 [*] `os` imported but unused
+  |
+3 | %matplotlib inline
+4 | 
+5 | import os
+  |        ^^ F401
+  |
+  = help: Remove unused import: `os`
+
+ℹ Fix
+2 2 | 
+3 3 | %matplotlib inline
+4 4 | 
+5   |-import os
+6 5 | 
+7 6 | _ = math.pi
+
+
diff --git a/crates/ruff/src/linter.rs b/crates/ruff/src/linter.rs
index 0f614535da42a..e4a12032e8121 100644
--- a/crates/ruff/src/linter.rs
+++ b/crates/ruff/src/linter.rs
@@ -7,7 +7,7 @@ use colored::Colorize;
 use itertools::Itertools;
 use log::error;
 use ruff_python_parser::lexer::LexResult;
-use ruff_python_parser::ParseError;
+use ruff_python_parser::{Mode, ParseError};
 use rustc_hash::FxHashMap;
 
 use ruff_diagnostics::Diagnostic;
@@ -138,7 +138,11 @@ pub fn check_path(
             .iter_enabled()
             .any(|rule_code| rule_code.lint_source().is_imports());
     if use_ast || use_imports || use_doc_lines {
-        match ruff_python_parser::parse_program_tokens(tokens, &path.to_string_lossy()) {
+        match ruff_python_parser::parse_program_tokens(
+            tokens,
+            &path.to_string_lossy(),
+            source_kind.map_or(false, SourceKind::is_jupyter),
+        ) {
             Ok(python_ast) => {
                 if use_ast {
                     diagnostics.extend(check_ast(
@@ -260,7 +264,7 @@ pub fn add_noqa_to_path(path: &Path, package: Option<&Path>, settings: &Settings
     let contents = std::fs::read_to_string(path)?;
 
     // Tokenize once.
-    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents);
+    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, Mode::Module);
 
     // Map row and column locations to byte slices (lazily).
     let locator = Locator::new(&contents);
@@ -327,8 +331,14 @@ pub fn lint_only(
     noqa: flags::Noqa,
     source_kind: Option<&SourceKind>,
 ) -> LinterResult<(Vec<Message>, Option<ImportMap>)> {
+    let mode = if source_kind.map_or(false, SourceKind::is_jupyter) {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+
     // Tokenize once.
-    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents);
+    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, mode);
 
     // Map row and column locations to byte slices (lazily).
     let locator = Locator::new(contents);
@@ -417,10 +427,16 @@ pub fn lint_fix<'a>(
     // Track whether the _initial_ source code was parseable.
     let mut parseable = false;
 
+    let mode = if source_kind.is_jupyter() {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+
     // Continuously autofix until the source code stabilizes.
     loop {
         // Tokenize once.
-        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&transformed);
+        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&transformed, mode);
 
         // Map row and column locations to byte slices (lazily).
         let locator = Locator::new(&transformed);
diff --git a/crates/ruff/src/rules/pyflakes/mod.rs b/crates/ruff/src/rules/pyflakes/mod.rs
index 7eb0434f31ef6..0d6c7d12eda1b 100644
--- a/crates/ruff/src/rules/pyflakes/mod.rs
+++ b/crates/ruff/src/rules/pyflakes/mod.rs
@@ -12,6 +12,7 @@ mod tests {
     use anyhow::Result;
     use regex::Regex;
     use ruff_python_parser::lexer::LexResult;
+    use ruff_python_parser::Mode;
     use test_case::test_case;
 
     use ruff_diagnostics::Diagnostic;
@@ -505,7 +506,7 @@ mod tests {
     fn flakes(contents: &str, expected: &[Rule]) {
         let contents = dedent(contents);
         let settings = Settings::for_rules(Linter::Pyflakes.rules());
-        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents);
+        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, Mode::Module);
         let locator = Locator::new(&contents);
         let stylist = Stylist::from_tokens(&tokens, &locator);
         let indexer = Indexer::from_tokens(&tokens, &locator);
diff --git a/crates/ruff/src/test.rs b/crates/ruff/src/test.rs
index 9e72487231e11..3dbdf89dccf10 100644
--- a/crates/ruff/src/test.rs
+++ b/crates/ruff/src/test.rs
@@ -7,6 +7,7 @@ use std::path::Path;
 use anyhow::Result;
 use itertools::Itertools;
 use ruff_python_parser::lexer::LexResult;
+use ruff_python_parser::Mode;
 use rustc_hash::FxHashMap;
 
 use ruff_diagnostics::{AutofixKind, Diagnostic};
@@ -99,8 +100,13 @@ pub(crate) fn max_iterations() -> usize {
 /// A convenient wrapper around [`check_path`], that additionally
 /// asserts that autofixes converge after a fixed number of iterations.
 fn test_contents(source_kind: &mut SourceKind, path: &Path, settings: &Settings) -> Vec<Message> {
+    let mode = if source_kind.is_jupyter() {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
     let contents = source_kind.content().to_string();
-    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents);
+    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, mode);
     let locator = Locator::new(&contents);
     let stylist = Stylist::from_tokens(&tokens, &locator);
     let indexer = Indexer::from_tokens(&tokens, &locator);
@@ -162,7 +168,7 @@ fn test_contents(source_kind: &mut SourceKind, path: &Path, settings: &Settings)
                 notebook.update(&source_map, &fixed_contents);
             };
 
-            let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&fixed_contents);
+            let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&fixed_contents, mode);
             let locator = Locator::new(&fixed_contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             let indexer = Indexer::from_tokens(&tokens, &locator);
diff --git a/crates/ruff_dev/src/print_ast.rs b/crates/ruff_dev/src/print_ast.rs
index fc141b13e7637..14aa14bfac6d7 100644
--- a/crates/ruff_dev/src/print_ast.rs
+++ b/crates/ruff_dev/src/print_ast.rs
@@ -5,18 +5,26 @@ use std::fs;
 use std::path::PathBuf;
 
 use anyhow::Result;
-use ruff_python_parser::parse_suite;
+use ruff_python_parser::{parse, Mode};
 
 #[derive(clap::Args)]
 pub(crate) struct Args {
     /// Python file for which to generate the AST.
     #[arg(required = true)]
     file: PathBuf,
+    /// Run in Jupyter mode i.e., allow line magics.
+    #[arg(long)]
+    jupyter: bool,
 }
 
 pub(crate) fn main(args: &Args) -> Result<()> {
     let contents = fs::read_to_string(&args.file)?;
-    let python_ast = parse_suite(&contents, &args.file.to_string_lossy())?;
+    let mode = if args.jupyter {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    let python_ast = parse(&contents, mode, &args.file.to_string_lossy())?;
     println!("{python_ast:#?}");
     Ok(())
 }
diff --git a/crates/ruff_dev/src/print_tokens.rs b/crates/ruff_dev/src/print_tokens.rs
index 573337e3f7fc2..c3e17904bcb05 100644
--- a/crates/ruff_dev/src/print_tokens.rs
+++ b/crates/ruff_dev/src/print_tokens.rs
@@ -12,11 +12,19 @@ pub(crate) struct Args {
     /// Python file for which to generate the AST.
     #[arg(required = true)]
     file: PathBuf,
+    /// Run in Jupyter mode i.e., allow line magics (`%`, `!`, `?`, `/`, `,`, `;`).
+    #[arg(long)]
+    jupyter: bool,
 }
 
 pub(crate) fn main(args: &Args) -> Result<()> {
     let contents = fs::read_to_string(&args.file)?;
-    for (tok, range) in lexer::lex(&contents, Mode::Module).flatten() {
+    let mode = if args.jupyter {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    for (tok, range) in lexer::lex(&contents, mode).flatten() {
         println!(
             "{start:#?} {tok:#?} {end:#?}",
             start = range.start(),
diff --git a/crates/ruff_python_parser/src/lib.rs b/crates/ruff_python_parser/src/lib.rs
index 0697099a43af2..a1b40bffda376 100644
--- a/crates/ruff_python_parser/src/lib.rs
+++ b/crates/ruff_python_parser/src/lib.rs
@@ -130,9 +130,9 @@ mod token;
 pub mod typing;
 
 /// Collect tokens up to and including the first error.
-pub fn tokenize(contents: &str) -> Vec<LexResult> {
+pub fn tokenize(contents: &str, mode: Mode) -> Vec<LexResult> {
     let mut tokens: Vec<LexResult> = vec![];
-    for tok in lexer::lex(contents, Mode::Module) {
+    for tok in lexer::lex(contents, mode) {
         let is_err = tok.is_err();
         tokens.push(tok);
         if is_err {
@@ -146,8 +146,14 @@ pub fn tokenize(contents: &str) -> Vec<LexResult> {
 pub fn parse_program_tokens(
     lxr: Vec<LexResult>,
     source_path: &str,
+    is_jupyter_notebook: bool,
 ) -> anyhow::Result<Suite, ParseError> {
-    match parse_tokens(lxr, Mode::Module, source_path)? {
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    match parse_tokens(lxr, mode, source_path)? {
         Mod::Module(m) => Ok(m.body),
         Mod::Expression(_) => unreachable!("Mode::Module doesn't return other variant"),
     }
diff --git a/crates/ruff_shrinking/src/main.rs b/crates/ruff_shrinking/src/main.rs
index 4c62cfa35534f..5efef3a728464 100644
--- a/crates/ruff_shrinking/src/main.rs
+++ b/crates/ruff_shrinking/src/main.rs
@@ -36,6 +36,7 @@ use regex::Regex;
 use ruff_python_ast::statement_visitor::{walk_body, walk_stmt, StatementVisitor};
 use ruff_python_ast::visitor::{walk_expr, Visitor};
 use ruff_python_ast::{Expr, Ranged, Stmt, Suite};
+use ruff_python_parser::Mode;
 use ruff_text_size::TextRange;
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
@@ -275,7 +276,7 @@ impl Strategy for StrategyRemoveToken {
         input: &'a str,
         _ast: &'a Suite,
     ) -> Result<Box<dyn ExactSizeStringIter + 'a>> {
-        let token_ranges: Vec<_> = ruff_python_parser::tokenize(input)
+        let token_ranges: Vec<_> = ruff_python_parser::tokenize(input, Mode::Module)
             .into_iter()
             // At this point we know we have valid python code
             .map(Result::unwrap)
@@ -320,9 +321,9 @@ fn minimization_step(
     pattern: &Regex,
     last_strategy_and_idx: Option<(&'static dyn Strategy, usize)>,
 ) -> Result<Option<(&'static dyn Strategy, usize, String)>> {
-    let tokens = ruff_python_parser::tokenize(input);
-    let ast =
-        ruff_python_parser::parse_program_tokens(tokens, "input.py").context("not valid python")?;
+    let tokens = ruff_python_parser::tokenize(input, Mode::Module);
+    let ast = ruff_python_parser::parse_program_tokens(tokens, "input.py", false)
+        .context("not valid python")?;
 
     // Try the last succeeding strategy first, skipping all that failed last time
     if let Some((last_strategy, last_idx)) = last_strategy_and_idx {
diff --git a/crates/ruff_wasm/src/lib.rs b/crates/ruff_wasm/src/lib.rs
index c3f48a9f286bd..e0e6fc28fecda 100644
--- a/crates/ruff_wasm/src/lib.rs
+++ b/crates/ruff_wasm/src/lib.rs
@@ -197,7 +197,7 @@ impl Workspace {
 
     pub fn check(&self, contents: &str) -> Result<JsValue, Error> {
         // Tokenize once.
-        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents);
+        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, Mode::Module);
 
         // Map row and column locations to byte slices (lazily).
         let locator = Locator::new(contents);

From 67183d62741fa62727eb84de57fe3453d1bb9d50 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 26 Jul 2023 17:59:44 +0530
Subject: [PATCH 02/18] Add magics to import sorting tests for Jupyter Notebook

---
 .../test/fixtures/jupyter/isort.ipynb         | 17 ++++++++
 .../fixtures/jupyter/isort_expected.ipynb     | 17 ++++++++
 ...yter__notebook__tests__import_sorting.snap | 41 +++++++++++++++++++
 3 files changed, 75 insertions(+)

diff --git a/crates/ruff/resources/test/fixtures/jupyter/isort.ipynb b/crates/ruff/resources/test/fixtures/jupyter/isort.ipynb
index aef5ff2e8b8aa..9572f7b25e8ad 100644
--- a/crates/ruff/resources/test/fixtures/jupyter/isort.ipynb
+++ b/crates/ruff/resources/test/fixtures/jupyter/isort.ipynb
@@ -25,6 +25,23 @@
     "def foo():\n",
     "    pass"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "16214f6f-bb32-4594-81be-79fb27c6ec92",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "import sys\n",
+    "\n",
+    "%matplotlib \\\n",
+    "    --inline\n",
+    "\n",
+    "import math\n",
+    "import abc"
+   ]
   }
  ],
  "metadata": {
diff --git a/crates/ruff/resources/test/fixtures/jupyter/isort_expected.ipynb b/crates/ruff/resources/test/fixtures/jupyter/isort_expected.ipynb
index 009c598e71672..5118aa36615e2 100644
--- a/crates/ruff/resources/test/fixtures/jupyter/isort_expected.ipynb
+++ b/crates/ruff/resources/test/fixtures/jupyter/isort_expected.ipynb
@@ -27,6 +27,23 @@
     "def foo():\n",
     "    pass"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d6c55c6-4a34-4662-914b-4ee11c9c24a5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "from pathlib import Path\n",
+    "\n",
+    "%matplotlib \\\n",
+    "    --inline\n",
+    "\n",
+    "import abc\n",
+    "import math"
+   ]
   }
  ],
  "metadata": {
diff --git a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
index 240556c37576f..a1e4c636f0ca6 100644
--- a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
+++ b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
@@ -47,4 +47,45 @@ isort.ipynb:cell 2:1:1: I001 [*] Import block is un-sorted or un-formatted
 7 9 | def foo():
 8 10 |     pass
 
+isort.ipynb:cell 2:6:1: I001 [*] Import block is un-sorted or un-formatted
+   |
+ 4 |   def foo():
+ 5 |       pass
+ 6 | / from pathlib import Path
+ 7 | | import sys
+ 8 | | 
+ 9 | | %matplotlib \
+   | |_^ I001
+10 |       --inline
+   |
+   = help: Organize imports
+
+ℹ Fix
+6  6  | # Newline should be added here
+7  7  | def foo():
+8  8  |     pass
+   9  |+import sys
+9  10 | from pathlib import Path
+10    |-import sys
+11 11 | 
+12 12 | %matplotlib \
+13 13 |     --inline
+
+isort.ipynb:cell 3:5:1: I001 [*] Import block is un-sorted or un-formatted
+  |
+3 |       --inline
+4 |   
+5 | / import math
+6 | | import abc
+  |
+  = help: Organize imports
+
+ℹ Fix
+12 12 | %matplotlib \
+13 13 |     --inline
+14 14 | 
+   15 |+import abc
+15 16 | import math
+16    |-import abc
+
 

From 9d14de721281f7b8311fefd347e37c4658782bfe Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 26 Jul 2023 18:07:01 +0530
Subject: [PATCH 03/18] Fix imports

---
 crates/ruff/src/jupyter/notebook.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/crates/ruff/src/jupyter/notebook.rs b/crates/ruff/src/jupyter/notebook.rs
index d234c7f17ff0c..2079b39353727 100644
--- a/crates/ruff/src/jupyter/notebook.rs
+++ b/crates/ruff/src/jupyter/notebook.rs
@@ -7,6 +7,7 @@ use std::path::Path;
 
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
+use rustpython_parser::lexer::lex;
 use rustpython_parser::Mode;
 use serde::Serialize;
 use serde_json::error::Category;

From 4c5dd4ebc44c99e6817a6d436a4adbfbb67f2e47 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 26 Jul 2023 22:28:30 +0530
Subject: [PATCH 04/18] Actually resolve merge conflicts

---
 crates/ruff/src/importer/insertion.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/crates/ruff/src/importer/insertion.rs b/crates/ruff/src/importer/insertion.rs
index 176fef854e1bd..6db795d0ba448 100644
--- a/crates/ruff/src/importer/insertion.rs
+++ b/crates/ruff/src/importer/insertion.rs
@@ -312,7 +312,7 @@ mod tests {
     fn start_of_file() -> Result<()> {
         fn insert(contents: &str) -> Result<Insertion> {
             let program = parse_suite(contents, "<filename>")?;
-            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents, Mode::Module);
+            let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             Ok(Insertion::start_of_file(&program, &locator, &stylist))
@@ -423,7 +423,7 @@ x = 1
     #[test]
     fn start_of_block() {
         fn insert(contents: &str, offset: TextSize) -> Insertion {
-            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents, Mode::Module);
+            let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             Insertion::start_of_block(offset, &locator, &stylist)

From 96dd1773747ff7af5d76a1078842d757a820cd29 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 26 Jul 2023 23:27:22 +0530
Subject: [PATCH 05/18] Use Jupyter mode for lexing

---
 crates/ruff/src/autofix/edits.rs              | 12 +++-
 crates/ruff/src/checkers/ast/mod.rs           |  5 +-
 crates/ruff/src/checkers/imports.rs           |  8 ++-
 crates/ruff/src/importer/insertion.rs         | 12 +++-
 crates/ruff/src/importer/mod.rs               | 13 +++-
 .../src/rules/flake8_annotations/fixes.rs     |  8 ++-
 .../flake8_annotations/rules/definition.rs    | 18 ++++--
 .../flake8_pytest_style/rules/fixture.rs      |  1 +
 .../flake8_pytest_style/rules/parametrize.rs  | 24 ++++---
 .../unnecessary_paren_on_raise_exception.rs   | 16 ++++-
 .../src/rules/flake8_simplify/rules/ast_if.rs |  1 +
 .../rules/flake8_simplify/rules/ast_with.rs   |  1 +
 .../rules/typing_only_runtime_import.rs       |  1 +
 crates/ruff/src/rules/isort/annotate.rs       |  3 +-
 crates/ruff/src/rules/isort/comments.rs       | 13 +++-
 crates/ruff/src/rules/isort/helpers.rs        | 13 +++-
 crates/ruff/src/rules/isort/mod.rs            |  9 ++-
 .../src/rules/isort/rules/organize_imports.rs |  3 +
 .../pandas_vet/rules/inplace_argument.rs      | 14 ++++-
 .../rules/f_string_missing_placeholders.rs    | 10 ++-
 .../rules/pyflakes/rules/unused_variable.rs   | 62 +++++++++++++------
 .../pylint/rules/bad_string_format_type.rs    |  7 ++-
 .../rules/printf_string_formatting.rs         | 13 ++--
 .../pyupgrade/rules/redundant_open_modes.rs   | 20 +++++-
 .../pyupgrade/rules/replace_stdout_stderr.rs  | 12 +++-
 .../rules/unnecessary_encode_utf8.rs          | 19 +++++-
 .../rules/useless_object_inheritance.rs       |  1 +
 crates/ruff_python_parser/src/lib.rs          | 14 ++++-
 28 files changed, 261 insertions(+), 72 deletions(-)

diff --git a/crates/ruff/src/autofix/edits.rs b/crates/ruff/src/autofix/edits.rs
index 5e3ff66e6c775..8c6c0863e4893 100644
--- a/crates/ruff/src/autofix/edits.rs
+++ b/crates/ruff/src/autofix/edits.rs
@@ -81,9 +81,15 @@ pub(crate) fn remove_argument(
     args: &[Expr],
     keywords: &[Keyword],
     remove_parentheses: bool,
+    is_jupyter_notebook: bool,
 ) -> Result<Edit> {
     // TODO(sbrugman): Preserve trailing comments.
     let contents = locator.after(call_at);
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
 
     let mut fix_start = None;
     let mut fix_end = None;
@@ -96,7 +102,7 @@ pub(crate) fn remove_argument(
     if n_arguments == 1 {
         // Case 1: there is only one argument.
         let mut count = 0u32;
-        for (tok, range) in lexer::lex_starts_at(contents, Mode::Module, call_at).flatten() {
+        for (tok, range) in lexer::lex_starts_at(contents, mode, call_at).flatten() {
             if tok.is_lpar() {
                 if count == 0 {
                     fix_start = Some(if remove_parentheses {
@@ -128,7 +134,7 @@ pub(crate) fn remove_argument(
     {
         // Case 2: argument or keyword is _not_ the last node.
         let mut seen_comma = false;
-        for (tok, range) in lexer::lex_starts_at(contents, Mode::Module, call_at).flatten() {
+        for (tok, range) in lexer::lex_starts_at(contents, mode, call_at).flatten() {
             if seen_comma {
                 if tok.is_non_logical_newline() {
                     // Also delete any non-logical newlines after the comma.
@@ -151,7 +157,7 @@ pub(crate) fn remove_argument(
     } else {
         // Case 3: argument or keyword is the last node, so we have to find the last
         // comma in the stmt.
-        for (tok, range) in lexer::lex_starts_at(contents, Mode::Module, call_at).flatten() {
+        for (tok, range) in lexer::lex_starts_at(contents, mode, call_at).flatten() {
             if range.start() == expr_range.start() {
                 fix_end = Some(expr_range.end());
                 break;
diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
index 41bddfdcbc65e..e23cc0d700ccd 100644
--- a/crates/ruff/src/checkers/ast/mod.rs
+++ b/crates/ruff/src/checkers/ast/mod.rs
@@ -52,7 +52,7 @@ use ruff_python_semantic::{
     ModuleKind, ScopeId, ScopeKind, SemanticModel, SemanticModelFlags, StarImport, SubmoduleImport,
 };
 use ruff_python_stdlib::builtins::{BUILTINS, MAGIC_GLOBALS};
-use ruff_python_stdlib::path::is_python_stub_file;
+use ruff_python_stdlib::path::{is_jupyter_notebook, is_python_stub_file};
 use ruff_source_file::Locator;
 
 use crate::checkers::ast::deferred::Deferred;
@@ -76,6 +76,8 @@ pub(crate) struct Checker<'a> {
     module_path: Option<&'a [String]>,
     /// Whether the current file is a stub (`.pyi`) file.
     is_stub: bool,
+    /// Whether the current file is a Jupyter notebook (`.ipynb`) file.
+    pub(crate) is_jupyter_notebook: bool,
     /// The [`flags::Noqa`] for the current analysis (i.e., whether to respect suppression
     /// comments).
     noqa: flags::Noqa,
@@ -126,6 +128,7 @@ impl<'a> Checker<'a> {
             package,
             module_path: module.path(),
             is_stub: is_python_stub_file(path),
+            is_jupyter_notebook: is_jupyter_notebook(path),
             locator,
             stylist,
             indexer,
diff --git a/crates/ruff/src/checkers/imports.rs b/crates/ruff/src/checkers/imports.rs
index 7eb17a2eca30d..7a16036445a51 100644
--- a/crates/ruff/src/checkers/imports.rs
+++ b/crates/ruff/src/checkers/imports.rs
@@ -104,7 +104,13 @@ pub(crate) fn check_imports(
         for block in &blocks {
             if !block.imports.is_empty() {
                 if let Some(diagnostic) = isort::rules::organize_imports(
-                    block, locator, stylist, indexer, settings, package,
+                    block,
+                    locator,
+                    stylist,
+                    indexer,
+                    settings,
+                    package,
+                    source_kind.map_or(false, SourceKind::is_jupyter),
                 ) {
                     diagnostics.push(diagnostic);
                 }
diff --git a/crates/ruff/src/importer/insertion.rs b/crates/ruff/src/importer/insertion.rs
index 6db795d0ba448..3df2781a57ab3 100644
--- a/crates/ruff/src/importer/insertion.rs
+++ b/crates/ruff/src/importer/insertion.rs
@@ -137,6 +137,7 @@ impl<'a> Insertion<'a> {
         mut location: TextSize,
         locator: &Locator<'a>,
         stylist: &Stylist,
+        is_jupyter_notebook: bool,
     ) -> Insertion<'a> {
         enum Awaiting {
             Colon(u32),
@@ -144,9 +145,14 @@ impl<'a> Insertion<'a> {
             Indent,
         }
 
+        let mode = if is_jupyter_notebook {
+            Mode::Jupyter
+        } else {
+            Mode::Module
+        };
+
         let mut state = Awaiting::Colon(0);
-        for (tok, range) in
-            lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten()
+        for (tok, range) in lexer::lex_starts_at(locator.after(location), mode, location).flatten()
         {
             match state {
                 // Iterate until we find the colon indicating the start of the block body.
@@ -426,7 +432,7 @@ x = 1
             let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
-            Insertion::start_of_block(offset, &locator, &stylist)
+            Insertion::start_of_block(offset, &locator, &stylist, false)
         }
 
         let contents = "if True: pass";
diff --git a/crates/ruff/src/importer/mod.rs b/crates/ruff/src/importer/mod.rs
index e08208237b72f..e3860f4fa7ad8 100644
--- a/crates/ruff/src/importer/mod.rs
+++ b/crates/ruff/src/importer/mod.rs
@@ -121,6 +121,7 @@ impl<'a> Importer<'a> {
         import: &StmtImports,
         at: TextSize,
         semantic: &SemanticModel,
+        is_jupyter_notebook: bool,
     ) -> Result<TypingImportEdit> {
         // Generate the modified import statement.
         let content = autofix::codemods::retain_imports(
@@ -140,7 +141,7 @@ impl<'a> Importer<'a> {
         // Add the import to a `TYPE_CHECKING` block.
         let add_import_edit = if let Some(block) = self.preceding_type_checking_block(at) {
             // Add the import to the `TYPE_CHECKING` block.
-            self.add_to_type_checking_block(&content, block.start())
+            self.add_to_type_checking_block(&content, block.start(), is_jupyter_notebook)
         } else {
             // Add the import to a new `TYPE_CHECKING` block.
             self.add_type_checking_block(
@@ -353,8 +354,14 @@ impl<'a> Importer<'a> {
     }
 
     /// Add an import statement to an existing `TYPE_CHECKING` block.
-    fn add_to_type_checking_block(&self, content: &str, at: TextSize) -> Edit {
-        Insertion::start_of_block(at, self.locator, self.stylist).into_edit(content)
+    fn add_to_type_checking_block(
+        &self,
+        content: &str,
+        at: TextSize,
+        is_jupyter_notebook: bool,
+    ) -> Edit {
+        Insertion::start_of_block(at, self.locator, self.stylist, is_jupyter_notebook)
+            .into_edit(content)
     }
 
     /// Return the import statement that precedes the given position, if any.
diff --git a/crates/ruff/src/rules/flake8_annotations/fixes.rs b/crates/ruff/src/rules/flake8_annotations/fixes.rs
index 2111978b0458e..2481c795126e8 100644
--- a/crates/ruff/src/rules/flake8_annotations/fixes.rs
+++ b/crates/ruff/src/rules/flake8_annotations/fixes.rs
@@ -10,14 +10,20 @@ pub(crate) fn add_return_annotation(
     locator: &Locator,
     stmt: &Stmt,
     annotation: &str,
+    is_jupyter_notebook: bool,
 ) -> Result<Edit> {
     let contents = &locator.contents()[stmt.range()];
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
 
     // Find the colon (following the `def` keyword).
     let mut seen_lpar = false;
     let mut seen_rpar = false;
     let mut count = 0u32;
-    for (tok, range) in lexer::lex_starts_at(contents, Mode::Module, stmt.start()).flatten() {
+    for (tok, range) in lexer::lex_starts_at(contents, mode, stmt.start()).flatten() {
         if seen_lpar && seen_rpar {
             if matches!(tok, Tok::Colon) {
                 return Ok(Edit::insertion(format!(" -> {annotation}"), range.start()));
diff --git a/crates/ruff/src/rules/flake8_annotations/rules/definition.rs b/crates/ruff/src/rules/flake8_annotations/rules/definition.rs
index 26fd2c9a7ee21..757ab9417ccc2 100644
--- a/crates/ruff/src/rules/flake8_annotations/rules/definition.rs
+++ b/crates/ruff/src/rules/flake8_annotations/rules/definition.rs
@@ -709,8 +709,13 @@ pub(crate) fn definition(
                     );
                     if checker.patch(diagnostic.kind.rule()) {
                         diagnostic.try_set_fix(|| {
-                            fixes::add_return_annotation(checker.locator(), stmt, "None")
-                                .map(Fix::suggested)
+                            fixes::add_return_annotation(
+                                checker.locator(),
+                                stmt,
+                                "None",
+                                checker.is_jupyter_notebook,
+                            )
+                            .map(Fix::suggested)
                         });
                     }
                     diagnostics.push(diagnostic);
@@ -727,8 +732,13 @@ pub(crate) fn definition(
                 if checker.patch(diagnostic.kind.rule()) {
                     if let Some(return_type) = simple_magic_return_type(name) {
                         diagnostic.try_set_fix(|| {
-                            fixes::add_return_annotation(checker.locator(), stmt, return_type)
-                                .map(Fix::suggested)
+                            fixes::add_return_annotation(
+                                checker.locator(),
+                                stmt,
+                                return_type,
+                                checker.is_jupyter_notebook,
+                            )
+                            .map(Fix::suggested)
                         });
                     }
                 }
diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs
index 66199052d8903..e8e8a356c97a2 100644
--- a/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs
+++ b/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs
@@ -521,6 +521,7 @@ fn check_fixture_decorator(checker: &mut Checker, func_name: &str, decorator: &D
                                     args,
                                     keywords,
                                     false,
+                                    checker.is_jupyter_notebook,
                                 )
                                 .map(Fix::suggested)
                             });
diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
index b1283c1b35a14..c16d75f92d594 100644
--- a/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
+++ b/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
@@ -95,18 +95,19 @@ fn elts_to_csv(elts: &[Expr], generator: Generator) -> Option<String> {
 /// ```
 ///
 /// This method assumes that the first argument is a string.
-fn get_parametrize_name_range(decorator: &Decorator, expr: &Expr, locator: &Locator) -> TextRange {
+fn get_parametrize_name_range(
+    decorator: &Decorator,
+    expr: &Expr,
+    locator: &Locator,
+    mode: Mode,
+) -> TextRange {
     let mut locations = Vec::new();
     let mut implicit_concat = None;
 
     // The parenthesis are not part of the AST, so we need to tokenize the
     // decorator to find them.
-    for (tok, range) in lexer::lex_starts_at(
-        locator.slice(decorator.range()),
-        Mode::Module,
-        decorator.start(),
-    )
-    .flatten()
+    for (tok, range) in
+        lexer::lex_starts_at(locator.slice(decorator.range()), mode, decorator.start()).flatten()
     {
         match tok {
             Tok::Lpar => locations.push(range.start()),
@@ -131,6 +132,11 @@ fn get_parametrize_name_range(decorator: &Decorator, expr: &Expr, locator: &Loca
 /// PT006
 fn check_names(checker: &mut Checker, decorator: &Decorator, expr: &Expr) {
     let names_type = checker.settings.flake8_pytest_style.parametrize_names_type;
+    let mode = if checker.is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
 
     match expr {
         Expr::Constant(ast::ExprConstant {
@@ -142,7 +148,7 @@ fn check_names(checker: &mut Checker, decorator: &Decorator, expr: &Expr) {
                 match names_type {
                     types::ParametrizeNameType::Tuple => {
                         let name_range =
-                            get_parametrize_name_range(decorator, expr, checker.locator());
+                            get_parametrize_name_range(decorator, expr, checker.locator(), mode);
                         let mut diagnostic = Diagnostic::new(
                             PytestParametrizeNamesWrongType {
                                 expected: names_type,
@@ -173,7 +179,7 @@ fn check_names(checker: &mut Checker, decorator: &Decorator, expr: &Expr) {
                     }
                     types::ParametrizeNameType::List => {
                         let name_range =
-                            get_parametrize_name_range(decorator, expr, checker.locator());
+                            get_parametrize_name_range(decorator, expr, checker.locator(), mode);
                         let mut diagnostic = Diagnostic::new(
                             PytestParametrizeNamesWrongType {
                                 expected: names_type,
diff --git a/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs b/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
index 23e92918534e6..c5b750f793e02 100644
--- a/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
+++ b/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
@@ -64,7 +64,7 @@ pub(crate) fn unnecessary_paren_on_raise_exception(checker: &mut Checker, expr:
                 return;
             }
 
-            let range = match_parens(func.end(), checker.locator())
+            let range = match_parens(func.end(), checker.locator(), checker.is_jupyter_notebook)
                 .expect("Expected call to include parentheses");
             let mut diagnostic = Diagnostic::new(UnnecessaryParenOnRaiseException, range);
             if checker.patch(diagnostic.kind.rule()) {
@@ -76,14 +76,24 @@ pub(crate) fn unnecessary_paren_on_raise_exception(checker: &mut Checker, expr:
 }
 
 /// Return the range of the first parenthesis pair after a given [`TextSize`].
-fn match_parens(start: TextSize, locator: &Locator) -> Option<TextRange> {
+fn match_parens(
+    start: TextSize,
+    locator: &Locator,
+    is_jupyter_notebook: bool,
+) -> Option<TextRange> {
     let contents = &locator.contents()[usize::from(start)..];
 
     let mut fix_start = None;
     let mut fix_end = None;
     let mut count = 0u32;
 
-    for (tok, range) in lexer::lex_starts_at(contents, Mode::Module, start).flatten() {
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+
+    for (tok, range) in lexer::lex_starts_at(contents, mode, start).flatten() {
         match tok {
             Tok::Lpar => {
                 if count == 0 {
diff --git a/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs b/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs
index 0bcf4783f9963..9d73345de3633 100644
--- a/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs
+++ b/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs
@@ -374,6 +374,7 @@ pub(crate) fn nested_if_statements(checker: &mut Checker, stmt_if: &StmtIf, pare
     let colon = first_colon_range(
         TextRange::new(test.end(), first_stmt.start()),
         checker.locator().contents(),
+        checker.is_jupyter_notebook,
     );
 
     // Check if the parent is already emitting a larger diagnostic including this if statement
diff --git a/crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs b/crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs
index 2a9ff71350d8c..26f409ab1d440 100644
--- a/crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs
+++ b/crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs
@@ -119,6 +119,7 @@ pub(crate) fn multiple_with_statements(
                 body.first().expect("Expected body to be non-empty").start(),
             ),
             checker.locator().contents(),
+            checker.is_jupyter_notebook,
         );
 
         let mut diagnostic = Diagnostic::new(
diff --git a/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs b/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs
index 0447b9ebdca7a..aa7a054001846 100644
--- a/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs
+++ b/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs
@@ -447,6 +447,7 @@ fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result
         },
         at,
         checker.semantic(),
+        checker.is_jupyter_notebook,
     )?;
 
     Ok(
diff --git a/crates/ruff/src/rules/isort/annotate.rs b/crates/ruff/src/rules/isort/annotate.rs
index 1aeca71811e63..351ca6418bc2c 100644
--- a/crates/ruff/src/rules/isort/annotate.rs
+++ b/crates/ruff/src/rules/isort/annotate.rs
@@ -13,6 +13,7 @@ pub(crate) fn annotate_imports<'a>(
     comments: Vec<Comment<'a>>,
     locator: &Locator,
     split_on_trailing_comma: bool,
+    is_jupyter_notebook: bool,
 ) -> Vec<AnnotatedImport<'a>> {
     let mut comments_iter = comments.into_iter().peekable();
 
@@ -119,7 +120,7 @@ pub(crate) fn annotate_imports<'a>(
                         names: aliases,
                         level: level.map(|level| level.to_u32()),
                         trailing_comma: if split_on_trailing_comma {
-                            trailing_comma(import, locator)
+                            trailing_comma(import, locator, is_jupyter_notebook)
                         } else {
                             TrailingComma::default()
                         },
diff --git a/crates/ruff/src/rules/isort/comments.rs b/crates/ruff/src/rules/isort/comments.rs
index b7963ff1ab244..d0b1564f86497 100644
--- a/crates/ruff/src/rules/isort/comments.rs
+++ b/crates/ruff/src/rules/isort/comments.rs
@@ -22,9 +22,18 @@ impl Comment<'_> {
 }
 
 /// Collect all comments in an import block.
-pub(crate) fn collect_comments<'a>(range: TextRange, locator: &'a Locator) -> Vec<Comment<'a>> {
+pub(crate) fn collect_comments<'a>(
+    range: TextRange,
+    locator: &'a Locator,
+    is_jupyter_notebook: bool,
+) -> Vec<Comment<'a>> {
     let contents = locator.slice(range);
-    lexer::lex_starts_at(contents, Mode::Module, range.start())
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    lexer::lex_starts_at(contents, mode, range.start())
         .flatten()
         .filter_map(|(tok, range)| {
             if let Tok::Comment(value) = tok {
diff --git a/crates/ruff/src/rules/isort/helpers.rs b/crates/ruff/src/rules/isort/helpers.rs
index adf185891db5d..89038f7070f28 100644
--- a/crates/ruff/src/rules/isort/helpers.rs
+++ b/crates/ruff/src/rules/isort/helpers.rs
@@ -8,11 +8,20 @@ use crate::rules::isort::types::TrailingComma;
 
 /// Return `true` if a `Stmt::ImportFrom` statement ends with a magic
 /// trailing comma.
-pub(super) fn trailing_comma(stmt: &Stmt, locator: &Locator) -> TrailingComma {
+pub(super) fn trailing_comma(
+    stmt: &Stmt,
+    locator: &Locator,
+    is_jupyter_notebook: bool,
+) -> TrailingComma {
     let contents = locator.slice(stmt.range());
     let mut count = 0u32;
     let mut trailing_comma = TrailingComma::Absent;
-    for (tok, _) in lexer::lex_starts_at(contents, Mode::Module, stmt.start()).flatten() {
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    for (tok, _) in lexer::lex_starts_at(contents, mode, stmt.start()).flatten() {
         if matches!(tok, Tok::Lpar) {
             count = count.saturating_add(1);
         }
diff --git a/crates/ruff/src/rules/isort/mod.rs b/crates/ruff/src/rules/isort/mod.rs
index 854b26728b386..b095b727aeb71 100644
--- a/crates/ruff/src/rules/isort/mod.rs
+++ b/crates/ruff/src/rules/isort/mod.rs
@@ -72,6 +72,7 @@ pub(crate) fn format_imports(
     stylist: &Stylist,
     src: &[PathBuf],
     package: Option<&Path>,
+    is_jupyter_notebook: bool,
     combine_as_imports: bool,
     force_single_line: bool,
     force_sort_within_sections: bool,
@@ -94,7 +95,13 @@ pub(crate) fn format_imports(
     section_order: &[ImportSection],
 ) -> String {
     let trailer = &block.trailer;
-    let block = annotate_imports(&block.imports, comments, locator, split_on_trailing_comma);
+    let block = annotate_imports(
+        &block.imports,
+        comments,
+        locator,
+        split_on_trailing_comma,
+        is_jupyter_notebook,
+    );
 
     // Normalize imports (i.e., deduplicate, aggregate `from` imports).
     let block = normalize_imports(
diff --git a/crates/ruff/src/rules/isort/rules/organize_imports.rs b/crates/ruff/src/rules/isort/rules/organize_imports.rs
index 02873311d91cf..1cb8182f966e8 100644
--- a/crates/ruff/src/rules/isort/rules/organize_imports.rs
+++ b/crates/ruff/src/rules/isort/rules/organize_imports.rs
@@ -87,6 +87,7 @@ pub(crate) fn organize_imports(
     indexer: &Indexer,
     settings: &Settings,
     package: Option<&Path>,
+    is_jupyter_notebook: bool,
 ) -> Option<Diagnostic> {
     let indentation = locator.slice(extract_indentation_range(&block.imports, locator));
     let indentation = leading_indentation(indentation);
@@ -105,6 +106,7 @@ pub(crate) fn organize_imports(
     let comments = comments::collect_comments(
         TextRange::new(range.start(), locator.full_line_end(range.end())),
         locator,
+        is_jupyter_notebook,
     );
 
     let trailing_line_end = if block.trailer.is_none() {
@@ -123,6 +125,7 @@ pub(crate) fn organize_imports(
         stylist,
         &settings.src,
         package,
+        is_jupyter_notebook,
         settings.isort.combine_as_imports,
         settings.isort.force_single_line,
         settings.isort.force_sort_within_sections,
diff --git a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
index 7deb387c8c062..1fb616740e7ef 100644
--- a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
+++ b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
@@ -107,6 +107,7 @@ pub(crate) fn inplace_argument(
                             keyword.range(),
                             args,
                             keywords,
+                            checker.is_jupyter_notebook,
                         ) {
                             diagnostic.set_fix(fix);
                         }
@@ -130,6 +131,7 @@ fn convert_inplace_argument_to_assignment(
     expr_range: TextRange,
     args: &[Expr],
     keywords: &[Keyword],
+    is_jupyter_notebook: bool,
 ) -> Option<Fix> {
     // Add the assignment.
     let call = expr.as_call_expr()?;
@@ -140,8 +142,16 @@ fn convert_inplace_argument_to_assignment(
     );
 
     // Remove the `inplace` argument.
-    let remove_argument =
-        remove_argument(locator, call.func.end(), expr_range, args, keywords, false).ok()?;
+    let remove_argument = remove_argument(
+        locator,
+        call.func.end(),
+        expr_range,
+        args,
+        keywords,
+        false,
+        is_jupyter_notebook,
+    )
+    .ok()?;
 
     Some(Fix::suggested_edits(insert_assignment, [remove_argument]))
 }
diff --git a/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs b/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
index 0a247f872aaaf..9cdec4a422113 100644
--- a/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
+++ b/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
@@ -52,9 +52,10 @@ impl AlwaysAutofixableViolation for FStringMissingPlaceholders {
 fn find_useless_f_strings<'a>(
     expr: &'a Expr,
     locator: &'a Locator,
+    mode: Mode,
 ) -> impl Iterator<Item = (TextRange, TextRange)> + 'a {
     let contents = locator.slice(expr.range());
-    lexer::lex_starts_at(contents, Mode::Module, expr.start())
+    lexer::lex_starts_at(contents, mode, expr.start())
         .flatten()
         .filter_map(|(tok, range)| match tok {
             Tok::String {
@@ -81,11 +82,16 @@ fn find_useless_f_strings<'a>(
 
 /// F541
 pub(crate) fn f_string_missing_placeholders(expr: &Expr, values: &[Expr], checker: &mut Checker) {
+    let mode = if checker.is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
     if !values
         .iter()
         .any(|value| matches!(value, Expr::FormattedValue(_)))
     {
-        for (prefix_range, tok_range) in find_useless_f_strings(expr, checker.locator()) {
+        for (prefix_range, tok_range) in find_useless_f_strings(expr, checker.locator(), mode) {
             let mut diagnostic = Diagnostic::new(FStringMissingPlaceholders, tok_range);
             if checker.patch(diagnostic.kind.rule()) {
                 diagnostic.set_fix(convert_f_string_to_regular_string(
diff --git a/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs b/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
index 50c17d1085af9..20544044641b8 100644
--- a/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
+++ b/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
@@ -62,12 +62,17 @@ impl Violation for UnusedVariable {
 }
 
 /// Return the [`TextRange`] of the token before the next match of the predicate
-fn match_token_before<F>(location: TextSize, locator: &Locator, f: F) -> Option<TextRange>
+fn match_token_before<F>(
+    location: TextSize,
+    locator: &Locator,
+    mode: Mode,
+    f: F,
+) -> Option<TextRange>
 where
     F: Fn(Tok) -> bool,
 {
     let contents = locator.after(location);
-    for ((_, range), (tok, _)) in lexer::lex_starts_at(contents, Mode::Module, location)
+    for ((_, range), (tok, _)) in lexer::lex_starts_at(contents, mode, location)
         .flatten()
         .tuple_windows()
     {
@@ -80,7 +85,12 @@ where
 
 /// Return the [`TextRange`] of the token after the next match of the predicate, skipping over
 /// any bracketed expressions.
-fn match_token_after<F>(location: TextSize, locator: &Locator, f: F) -> Option<TextRange>
+fn match_token_after<F>(
+    location: TextSize,
+    locator: &Locator,
+    mode: Mode,
+    f: F,
+) -> Option<TextRange>
 where
     F: Fn(Tok) -> bool,
 {
@@ -91,7 +101,7 @@ where
     let mut sqb_count = 0u32;
     let mut brace_count = 0u32;
 
-    for ((tok, _), (_, range)) in lexer::lex_starts_at(contents, Mode::Module, location)
+    for ((tok, _), (_, range)) in lexer::lex_starts_at(contents, mode, location)
         .flatten()
         .tuple_windows()
     {
@@ -131,7 +141,12 @@ where
 
 /// Return the [`TextRange`] of the token matching the predicate or the first mismatched
 /// bracket, skipping over any bracketed expressions.
-fn match_token_or_closing_brace<F>(location: TextSize, locator: &Locator, f: F) -> Option<TextRange>
+fn match_token_or_closing_brace<F>(
+    location: TextSize,
+    locator: &Locator,
+    mode: Mode,
+    f: F,
+) -> Option<TextRange>
 where
     F: Fn(Tok) -> bool,
 {
@@ -142,7 +157,7 @@ where
     let mut sqb_count = 0u32;
     let mut brace_count = 0u32;
 
-    for (tok, range) in lexer::lex_starts_at(contents, Mode::Module, location).flatten() {
+    for (tok, range) in lexer::lex_starts_at(contents, mode, location).flatten() {
         match tok {
             Tok::Lpar => {
                 par_count = par_count.saturating_add(1);
@@ -194,6 +209,12 @@ fn remove_unused_variable(
     range: TextRange,
     checker: &Checker,
 ) -> Option<Fix> {
+    let mode = if checker.is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+
     // First case: simple assignment (`x = 1`)
     if let Stmt::Assign(ast::StmtAssign { targets, value, .. }) = stmt {
         if let Some(target) = targets.iter().find(|target| range == target.range()) {
@@ -204,8 +225,9 @@ fn remove_unused_variable(
                     // If the expression is complex (`x = foo()`), remove the assignment,
                     // but preserve the right-hand side.
                     let start = target.start();
-                    let end = match_token_after(start, checker.locator(), |tok| tok == Tok::Equal)?
-                        .start();
+                    let end =
+                        match_token_after(start, checker.locator(), mode, |tok| tok == Tok::Equal)?
+                            .start();
                     let edit = Edit::deletion(start, end);
                     Some(Fix::suggested(edit))
                 } else {
@@ -230,7 +252,8 @@ fn remove_unused_variable(
                 // but preserve the right-hand side.
                 let start = stmt.start();
                 let end =
-                    match_token_after(start, checker.locator(), |tok| tok == Tok::Equal)?.start();
+                    match_token_after(start, checker.locator(), mode, |tok| tok == Tok::Equal)?
+                        .start();
                 let edit = Edit::deletion(start, end);
                 Some(Fix::suggested(edit))
             } else {
@@ -249,17 +272,20 @@ fn remove_unused_variable(
             if let Some(optional_vars) = &item.optional_vars {
                 if optional_vars.range() == range {
                     // Find the first token before the `as` keyword.
-                    let start =
-                        match_token_before(item.context_expr.start(), checker.locator(), |tok| {
-                            tok == Tok::As
-                        })?
-                        .end();
+                    let start = match_token_before(
+                        item.context_expr.start(),
+                        checker.locator(),
+                        mode,
+                        |tok| tok == Tok::As,
+                    )?
+                    .end();
 
                     // Find the first colon, comma, or closing bracket after the `as` keyword.
-                    let end = match_token_or_closing_brace(start, checker.locator(), |tok| {
-                        tok == Tok::Colon || tok == Tok::Comma
-                    })?
-                    .start();
+                    let end =
+                        match_token_or_closing_brace(start, checker.locator(), mode, |tok| {
+                            tok == Tok::Colon || tok == Tok::Comma
+                        })?
+                        .start();
 
                     let edit = Edit::deletion(start, end);
                     return Some(Fix::suggested(edit));
diff --git a/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs b/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
index 857fc4e5db2b6..f49cf2aef3f7b 100644
--- a/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
+++ b/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
@@ -203,7 +203,12 @@ pub(crate) fn bad_string_format_type(checker: &mut Checker, expr: &Expr, right:
     // Grab each string segment (in case there's an implicit concatenation).
     let content = checker.locator().slice(expr.range());
     let mut strings: Vec<TextRange> = vec![];
-    for (tok, range) in lexer::lex_starts_at(content, Mode::Module, expr.start()).flatten() {
+    let mode = if checker.is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    for (tok, range) in lexer::lex_starts_at(content, mode, expr.start()).flatten() {
         if tok.is_string() {
             strings.push(range);
         } else if tok.is_percent() {
diff --git a/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs b/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
index 073639447fac0..6f9fb997dec46 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
@@ -337,12 +337,13 @@ pub(crate) fn printf_string_formatting(
     // Grab each string segment (in case there's an implicit concatenation).
     let mut strings: Vec<TextRange> = vec![];
     let mut extension = None;
-    for (tok, range) in lexer::lex_starts_at(
-        checker.locator().slice(expr.range()),
-        Mode::Module,
-        expr.start(),
-    )
-    .flatten()
+    let mode = if checker.is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    for (tok, range) in
+        lexer::lex_starts_at(checker.locator().slice(expr.range()), mode, expr.start()).flatten()
     {
         if tok.is_string() {
             strings.push(range);
diff --git a/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs b/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
index df0a1156b5bb6..bb7ba78a74367 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
@@ -84,6 +84,7 @@ pub(crate) fn redundant_open_modes(checker: &mut Checker, expr: &Expr) {
                                 mode.replacement_value(),
                                 checker.locator(),
                                 checker.patch(Rule::RedundantOpenModes),
+                                checker.is_jupyter_notebook,
                             ));
                         }
                     }
@@ -103,6 +104,7 @@ pub(crate) fn redundant_open_modes(checker: &mut Checker, expr: &Expr) {
                         mode.replacement_value(),
                         checker.locator(),
                         checker.patch(Rule::RedundantOpenModes),
+                        checker.is_jupyter_notebook,
                     ));
                 }
             }
@@ -182,6 +184,7 @@ fn create_check(
     replacement_value: Option<&str>,
     locator: &Locator,
     patch: bool,
+    is_jupyter_notebook: bool,
 ) -> Diagnostic {
     let mut diagnostic = Diagnostic::new(
         RedundantOpenModes {
@@ -197,14 +200,20 @@ fn create_check(
             )));
         } else {
             diagnostic.try_set_fix(|| {
-                create_remove_param_fix(locator, expr, mode_param).map(Fix::automatic)
+                create_remove_param_fix(locator, expr, mode_param, is_jupyter_notebook)
+                    .map(Fix::automatic)
             });
         }
     }
     diagnostic
 }
 
-fn create_remove_param_fix(locator: &Locator, expr: &Expr, mode_param: &Expr) -> Result<Edit> {
+fn create_remove_param_fix(
+    locator: &Locator,
+    expr: &Expr,
+    mode_param: &Expr,
+    is_jupyter_notebook: bool,
+) -> Result<Edit> {
     let content = locator.slice(expr.range());
     // Find the last comma before mode_param and create a deletion fix
     // starting from the comma and ending after mode_param.
@@ -212,7 +221,12 @@ fn create_remove_param_fix(locator: &Locator, expr: &Expr, mode_param: &Expr) ->
     let mut fix_end: Option<TextSize> = None;
     let mut is_first_arg: bool = false;
     let mut delete_first_arg: bool = false;
-    for (tok, range) in lexer::lex_starts_at(content, Mode::Module, expr.start()).flatten() {
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    for (tok, range) in lexer::lex_starts_at(content, mode, expr.start()).flatten() {
         if range.start() == mode_param.start() {
             if is_first_arg {
                 delete_first_arg = true;
diff --git a/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs b/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
index f52a3a70c4a34..676a390d26506 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
@@ -59,6 +59,7 @@ fn generate_fix(
     keywords: &[Keyword],
     stdout: &Keyword,
     stderr: &Keyword,
+    is_jupyter_notebook: bool,
 ) -> Result<Fix> {
     let (first, second) = if stdout.start() < stderr.start() {
         (stdout, stderr)
@@ -74,6 +75,7 @@ fn generate_fix(
             args,
             keywords,
             false,
+            is_jupyter_notebook,
         )?],
     ))
 }
@@ -115,7 +117,15 @@ pub(crate) fn replace_stdout_stderr(
         let mut diagnostic = Diagnostic::new(ReplaceStdoutStderr, expr.range());
         if checker.patch(diagnostic.kind.rule()) {
             diagnostic.try_set_fix(|| {
-                generate_fix(checker.locator(), func, args, keywords, stdout, stderr)
+                generate_fix(
+                    checker.locator(),
+                    func,
+                    args,
+                    keywords,
+                    stdout,
+                    stderr,
+                    checker.is_jupyter_notebook,
+                )
             });
         }
         checker.diagnostics.push(diagnostic);
diff --git a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
index 3f1b1dbdd2351..658a9fe6b61a2 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
@@ -122,12 +122,17 @@ fn match_encoding_arg<'a>(args: &'a [Expr], kwargs: &'a [Keyword]) -> Option<Enc
 }
 
 /// Return a [`Fix`] replacing the call to encode with a byte string.
-fn replace_with_bytes_literal(locator: &Locator, expr: &Expr) -> Fix {
+fn replace_with_bytes_literal(locator: &Locator, expr: &Expr, is_jupyter_notebook: bool) -> Fix {
     // Build up a replacement string by prefixing all string tokens with `b`.
     let contents = locator.slice(expr.range());
     let mut replacement = String::with_capacity(contents.len() + 1);
     let mut prev = expr.start();
-    for (tok, range) in lexer::lex_starts_at(contents, Mode::Module, expr.start()).flatten() {
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    for (tok, range) in lexer::lex_starts_at(contents, mode, expr.start()).flatten() {
         match tok {
             Tok::Dot => break,
             Tok::String { .. } => {
@@ -175,7 +180,11 @@ pub(crate) fn unnecessary_encode_utf8(
                         expr.range(),
                     );
                     if checker.patch(Rule::UnnecessaryEncodeUTF8) {
-                        diagnostic.set_fix(replace_with_bytes_literal(checker.locator(), expr));
+                        diagnostic.set_fix(replace_with_bytes_literal(
+                            checker.locator(),
+                            expr,
+                            checker.is_jupyter_notebook,
+                        ));
                     }
                     checker.diagnostics.push(diagnostic);
                 } else if let EncodingArg::Keyword(kwarg) = encoding_arg {
@@ -196,6 +205,7 @@ pub(crate) fn unnecessary_encode_utf8(
                                 args,
                                 kwargs,
                                 false,
+                                checker.is_jupyter_notebook,
                             )
                             .map(Fix::automatic)
                         });
@@ -218,6 +228,7 @@ pub(crate) fn unnecessary_encode_utf8(
                                 args,
                                 kwargs,
                                 false,
+                                checker.is_jupyter_notebook,
                             )
                             .map(Fix::automatic)
                         });
@@ -247,6 +258,7 @@ pub(crate) fn unnecessary_encode_utf8(
                                 args,
                                 kwargs,
                                 false,
+                                checker.is_jupyter_notebook,
                             )
                             .map(Fix::automatic)
                         });
@@ -269,6 +281,7 @@ pub(crate) fn unnecessary_encode_utf8(
                                 args,
                                 kwargs,
                                 false,
+                                checker.is_jupyter_notebook,
                             )
                             .map(Fix::automatic)
                         });
diff --git a/crates/ruff/src/rules/pyupgrade/rules/useless_object_inheritance.rs b/crates/ruff/src/rules/pyupgrade/rules/useless_object_inheritance.rs
index 32db85e693fdb..82a0e60985362 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/useless_object_inheritance.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/useless_object_inheritance.rs
@@ -73,6 +73,7 @@ pub(crate) fn useless_object_inheritance(checker: &mut Checker, class_def: &ast:
                     &class_def.bases,
                     &class_def.keywords,
                     true,
+                    checker.is_jupyter_notebook,
                 )?;
                 Ok(Fix::automatic(edit))
             });
diff --git a/crates/ruff_python_parser/src/lib.rs b/crates/ruff_python_parser/src/lib.rs
index a1b40bffda376..c30fa66d79236 100644
--- a/crates/ruff_python_parser/src/lib.rs
+++ b/crates/ruff_python_parser/src/lib.rs
@@ -160,9 +160,18 @@ pub fn parse_program_tokens(
 }
 
 /// Return the `Range` of the first `Tok::Colon` token in a `Range`.
-pub fn first_colon_range(range: TextRange, source: &str) -> Option<TextRange> {
+pub fn first_colon_range(
+    range: TextRange,
+    source: &str,
+    is_jupyter_notebook: bool,
+) -> Option<TextRange> {
     let contents = &source[range];
-    let range = lexer::lex_starts_at(contents, Mode::Module, range.start())
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    let range = lexer::lex_starts_at(contents, mode, range.start())
         .flatten()
         .find(|(tok, _)| tok.is_colon())
         .map(|(_, range)| range);
@@ -363,6 +372,7 @@ mod tests {
         let range = first_colon_range(
             TextRange::new(TextSize::from(0), contents.text_len()),
             contents,
+            false,
         )
         .unwrap();
         assert_eq!(&contents[range], ":");

From adb790a391e23e1b517c5d877133ca01df89cddc Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Thu, 27 Jul 2023 07:11:33 +0530
Subject: [PATCH 06/18] Use `is_jupyter_notebook` param

---
 .../flake8_pytest_style/rules/parametrize.rs  | 29 ++++++++++++-------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
index c16d75f92d594..01696bb3ecc85 100644
--- a/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
+++ b/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
@@ -99,11 +99,17 @@ fn get_parametrize_name_range(
     decorator: &Decorator,
     expr: &Expr,
     locator: &Locator,
-    mode: Mode,
+    is_jupyter_notebook: bool,
 ) -> TextRange {
     let mut locations = Vec::new();
     let mut implicit_concat = None;
 
+    let mode = if is_jupyter_notebook {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+
     // The parenthesis are not part of the AST, so we need to tokenize the
     // decorator to find them.
     for (tok, range) in
@@ -132,11 +138,6 @@ fn get_parametrize_name_range(
 /// PT006
 fn check_names(checker: &mut Checker, decorator: &Decorator, expr: &Expr) {
     let names_type = checker.settings.flake8_pytest_style.parametrize_names_type;
-    let mode = if checker.is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
 
     match expr {
         Expr::Constant(ast::ExprConstant {
@@ -147,8 +148,12 @@ fn check_names(checker: &mut Checker, decorator: &Decorator, expr: &Expr) {
             if names.len() > 1 {
                 match names_type {
                     types::ParametrizeNameType::Tuple => {
-                        let name_range =
-                            get_parametrize_name_range(decorator, expr, checker.locator(), mode);
+                        let name_range = get_parametrize_name_range(
+                            decorator,
+                            expr,
+                            checker.locator(),
+                            checker.is_jupyter_notebook,
+                        );
                         let mut diagnostic = Diagnostic::new(
                             PytestParametrizeNamesWrongType {
                                 expected: names_type,
@@ -178,8 +183,12 @@ fn check_names(checker: &mut Checker, decorator: &Decorator, expr: &Expr) {
                         checker.diagnostics.push(diagnostic);
                     }
                     types::ParametrizeNameType::List => {
-                        let name_range =
-                            get_parametrize_name_range(decorator, expr, checker.locator(), mode);
+                        let name_range = get_parametrize_name_range(
+                            decorator,
+                            expr,
+                            checker.locator(),
+                            checker.is_jupyter_notebook,
+                        );
                         let mut diagnostic = Diagnostic::new(
                             PytestParametrizeNamesWrongType {
                                 expected: names_type,

From 0b958d62d63d0d6e218e33f686acc76e72b8408d Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Thu, 27 Jul 2023 07:58:36 +0530
Subject: [PATCH 07/18] Return both original and updated `SourceKind` in tests

---
 crates/ruff/src/jupyter/notebook.rs           |  6 ++--
 ...yter__notebook__tests__import_sorting.snap | 32 +++++++++----------
 ...jupyter__notebook__tests__line_magics.snap |  2 ++
 crates/ruff/src/test.rs                       |  5 +--
 4 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/crates/ruff/src/jupyter/notebook.rs b/crates/ruff/src/jupyter/notebook.rs
index 2079b39353727..abd13b590ca93 100644
--- a/crates/ruff/src/jupyter/notebook.rs
+++ b/crates/ruff/src/jupyter/notebook.rs
@@ -563,7 +563,7 @@ print("after empty cells")
     #[test]
     fn test_import_sorting() -> Result<()> {
         let path = "isort.ipynb".to_string();
-        let (diagnostics, source_kind) = test_notebook_path(
+        let (diagnostics, source_kind, _) = test_notebook_path(
             &path,
             Path::new("isort_expected.ipynb"),
             &settings::Settings::for_rule(Rule::UnsortedImports),
@@ -575,7 +575,7 @@ print("after empty cells")
     #[test]
     fn test_line_magics() -> Result<()> {
         let path = "line_magics.ipynb".to_string();
-        let (diagnostics, source_kind) = test_notebook_path(
+        let (diagnostics, source_kind, _) = test_notebook_path(
             &path,
             Path::new("line_magics_expected.ipynb"),
             &settings::Settings::for_rule(Rule::UnusedImport),
@@ -587,7 +587,7 @@ print("after empty cells")
     #[test]
     fn test_json_consistency() -> Result<()> {
         let path = "before_fix.ipynb".to_string();
-        let (_, source_kind) = test_notebook_path(
+        let (_, _, source_kind) = test_notebook_path(
             path,
             Path::new("after_fix.ipynb"),
             &settings::Settings::for_rule(Rule::UnusedImport),
diff --git a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
index a1e4c636f0ca6..e0ab3572a84d4 100644
--- a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
+++ b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
@@ -47,18 +47,16 @@ isort.ipynb:cell 2:1:1: I001 [*] Import block is un-sorted or un-formatted
 7 9 | def foo():
 8 10 |     pass
 
-isort.ipynb:cell 2:6:1: I001 [*] Import block is un-sorted or un-formatted
-   |
- 4 |   def foo():
- 5 |       pass
- 6 | / from pathlib import Path
- 7 | | import sys
- 8 | | 
- 9 | | %matplotlib \
-   | |_^ I001
-10 |       --inline
-   |
-   = help: Organize imports
+isort.ipynb:cell 3:1:1: I001 [*] Import block is un-sorted or un-formatted
+  |
+1 | / from pathlib import Path
+2 | | import sys
+3 | | 
+4 | | %matplotlib \
+  | |_^ I001
+5 |       --inline
+  |
+  = help: Organize imports
 
 ℹ Fix
 6  6  | # Newline should be added here
@@ -71,12 +69,12 @@ isort.ipynb:cell 2:6:1: I001 [*] Import block is un-sorted or un-formatted
 12 12 | %matplotlib \
 13 13 |     --inline
 
-isort.ipynb:cell 3:5:1: I001 [*] Import block is un-sorted or un-formatted
+isort.ipynb:cell 3:7:1: I001 [*] Import block is un-sorted or un-formatted
   |
-3 |       --inline
-4 |   
-5 | / import math
-6 | | import abc
+5 |       --inline
+6 |   
+7 | / import math
+8 | | import abc
   |
   = help: Organize imports
 
diff --git a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap
index 576b938b6b2e0..61211b2c18edb 100644
--- a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap
+++ b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap
@@ -7,6 +7,8 @@ line_magics.ipynb:cell 1:5:8: F401 [*] `os` imported but unused
 4 | 
 5 | import os
   |        ^^ F401
+6 | 
+7 | _ = math.pi
   |
   = help: Remove unused import: `os`
 
diff --git a/crates/ruff/src/test.rs b/crates/ruff/src/test.rs
index 3dbdf89dccf10..f8495064039e6 100644
--- a/crates/ruff/src/test.rs
+++ b/crates/ruff/src/test.rs
@@ -62,8 +62,9 @@ pub(crate) fn test_notebook_path(
     path: impl AsRef<Path>,
     expected: impl AsRef<Path>,
     settings: &Settings,
-) -> Result<(Vec<Message>, SourceKind)> {
+) -> Result<(Vec<Message>, SourceKind, SourceKind)> {
     let mut source_kind = SourceKind::Jupyter(read_jupyter_notebook(path.as_ref())?);
+    let original_source_kind = source_kind.clone();
     let messages = test_contents(&mut source_kind, path.as_ref(), settings);
     let expected_notebook = read_jupyter_notebook(expected.as_ref())?;
     if let SourceKind::Jupyter(notebook) = &source_kind {
@@ -71,7 +72,7 @@ pub(crate) fn test_notebook_path(
         assert_eq!(notebook.index(), expected_notebook.index());
         assert_eq!(notebook.content(), expected_notebook.content());
     };
-    Ok((messages, source_kind))
+    Ok((messages, original_source_kind, source_kind))
 }
 
 /// Run [`check_path`] on a snippet of Python code.

From 2e953cd50f50f5e72d2c333c578857eb2bfccb16 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Thu, 27 Jul 2023 10:24:21 +0530
Subject: [PATCH 08/18] Update test case to include cell magic

---
 .../resources/test/fixtures/jupyter/line_magics.ipynb | 11 +++++++++++
 .../test/fixtures/jupyter/line_magics_expected.ipynb  | 11 +++++++++++
 2 files changed, 22 insertions(+)

diff --git a/crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb b/crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb
index 785a73102ad73..5e9b10bb7b0e1 100644
--- a/crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb
+++ b/crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb
@@ -15,6 +15,17 @@
     "\n",
     "_ = math.pi"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2b0e2986-1b87-4bb6-9b1d-c11ca1decd87",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%timeit\n",
+    "import sys"
+   ]
   }
  ],
  "metadata": {
diff --git a/crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb b/crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb
index cdf69fa719600..8419f031e78f8 100644
--- a/crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb
+++ b/crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb
@@ -14,6 +14,17 @@
     "\n",
     "_ = math.pi"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d7b8e967-8b4a-493b-b6f7-d5cecfb3a5c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%timeit\n",
+    "import sys"
+   ]
   }
  ],
  "metadata": {

From 9bc3764d7169493c25dca92ddcd63aa642fa1086 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Fri, 28 Jul 2023 06:36:26 +0530
Subject: [PATCH 09/18] Remove leftover imports from rebase

---
 crates/ruff/src/jupyter/notebook.rs | 2 --
 1 file changed, 2 deletions(-)

diff --git a/crates/ruff/src/jupyter/notebook.rs b/crates/ruff/src/jupyter/notebook.rs
index abd13b590ca93..652e54850f05f 100644
--- a/crates/ruff/src/jupyter/notebook.rs
+++ b/crates/ruff/src/jupyter/notebook.rs
@@ -7,8 +7,6 @@ use std::path::Path;
 
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
-use rustpython_parser::lexer::lex;
-use rustpython_parser::Mode;
 use serde::Serialize;
 use serde_json::error::Category;
 

From 0d27e56813d261c31ab57149cd58592a939b6e20 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Fri, 28 Jul 2023 15:09:51 +0530
Subject: [PATCH 10/18] Use `PySourceType` for detecting lexing/parsing mode

---
 crates/ruff/src/autofix/edits.rs              | 19 ++++----
 .../ruff/src/checkers/ast/analyze/bindings.rs |  2 +-
 .../checkers/ast/analyze/deferred_scopes.rs   |  4 +-
 .../src/checkers/ast/analyze/definitions.rs   |  2 +-
 .../src/checkers/ast/analyze/expression.rs    | 30 ++++++------
 .../src/checkers/ast/analyze/parameters.rs    |  2 +-
 .../src/checkers/ast/analyze/statement.rs     | 20 ++++----
 crates/ruff/src/checkers/ast/mod.rs           | 16 +++----
 crates/ruff/src/checkers/imports.rs           | 18 ++++---
 crates/ruff/src/importer/insertion.rs         | 19 ++++----
 crates/ruff/src/importer/mod.rs               | 10 ++--
 crates/ruff/src/linter.rs                     | 45 +++++++++---------
 .../src/rules/flake8_annotations/fixes.rs     | 15 +++---
 .../flake8_annotations/rules/definition.rs    |  4 +-
 .../flake8_pytest_style/rules/fixture.rs      |  2 +-
 .../flake8_pytest_style/rules/parametrize.rs  | 23 +++++----
 .../unnecessary_paren_on_raise_exception.rs   | 15 ++----
 .../src/rules/flake8_simplify/rules/ast_if.rs |  2 +-
 .../rules/flake8_simplify/rules/ast_with.rs   |  2 +-
 .../rules/typing_only_runtime_import.rs       |  2 +-
 crates/ruff/src/rules/isort/annotate.rs       |  6 ++-
 crates/ruff/src/rules/isort/comments.rs       | 13 ++---
 crates/ruff/src/rules/isort/helpers.rs        | 12 ++---
 crates/ruff/src/rules/isort/mod.rs            |  5 +-
 .../rules/isort/rules/add_required_imports.rs | 11 +++--
 .../src/rules/isort/rules/organize_imports.rs |  7 +--
 .../pandas_vet/rules/inplace_argument.rs      |  7 +--
 crates/ruff/src/rules/pyflakes/mod.rs         |  7 ++-
 .../rules/f_string_missing_placeholders.rs    | 16 +++----
 .../rules/pyflakes/rules/unused_variable.rs   | 47 ++++++++++---------
 .../pylint/rules/bad_string_format_type.rs    | 11 ++---
 .../rules/printf_string_formatting.rs         | 15 +++---
 .../pyupgrade/rules/redundant_open_modes.rs   | 22 ++++-----
 .../pyupgrade/rules/replace_stdout_stderr.rs  |  7 +--
 .../rules/unnecessary_encode_utf8.rs          | 24 +++++-----
 .../rules/useless_object_inheritance.rs       |  2 +-
 crates/ruff/src/source_kind.rs                | 43 +++++++++++++++++
 crates/ruff/src/test.rs                       | 17 ++++---
 crates/ruff_benchmark/benches/linter.rs       |  5 +-
 crates/ruff_cli/src/diagnostics.rs            | 16 +++++--
 crates/ruff_wasm/src/lib.rs                   |  6 ++-
 41 files changed, 300 insertions(+), 251 deletions(-)

diff --git a/crates/ruff/src/autofix/edits.rs b/crates/ruff/src/autofix/edits.rs
index 8c6c0863e4893..1761528fe986d 100644
--- a/crates/ruff/src/autofix/edits.rs
+++ b/crates/ruff/src/autofix/edits.rs
@@ -1,7 +1,7 @@
 //! Interface for generating autofix edits from higher-level actions (e.g., "remove an argument").
 use anyhow::{bail, Result};
 use ruff_python_ast::{self as ast, ExceptHandler, Expr, Keyword, Ranged, Stmt};
-use ruff_python_parser::{lexer, Mode};
+use ruff_python_parser::lexer;
 use ruff_text_size::{TextLen, TextRange, TextSize};
 
 use ruff_diagnostics::Edit;
@@ -11,6 +11,7 @@ use ruff_python_trivia::{has_leading_content, is_python_whitespace, PythonWhites
 use ruff_source_file::{Locator, NewlineWithTrailingNewline};
 
 use crate::autofix::codemods;
+use crate::source_kind::PySourceType;
 
 /// Return the `Fix` to use when deleting a `Stmt`.
 ///
@@ -81,15 +82,10 @@ pub(crate) fn remove_argument(
     args: &[Expr],
     keywords: &[Keyword],
     remove_parentheses: bool,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Result<Edit> {
     // TODO(sbrugman): Preserve trailing comments.
     let contents = locator.after(call_at);
-    let mode = if is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
 
     let mut fix_start = None;
     let mut fix_end = None;
@@ -102,7 +98,8 @@ pub(crate) fn remove_argument(
     if n_arguments == 1 {
         // Case 1: there is only one argument.
         let mut count = 0u32;
-        for (tok, range) in lexer::lex_starts_at(contents, mode, call_at).flatten() {
+        for (tok, range) in lexer::lex_starts_at(contents, source_type.as_mode(), call_at).flatten()
+        {
             if tok.is_lpar() {
                 if count == 0 {
                     fix_start = Some(if remove_parentheses {
@@ -134,7 +131,8 @@ pub(crate) fn remove_argument(
     {
         // Case 2: argument or keyword is _not_ the last node.
         let mut seen_comma = false;
-        for (tok, range) in lexer::lex_starts_at(contents, mode, call_at).flatten() {
+        for (tok, range) in lexer::lex_starts_at(contents, source_type.as_mode(), call_at).flatten()
+        {
             if seen_comma {
                 if tok.is_non_logical_newline() {
                     // Also delete any non-logical newlines after the comma.
@@ -157,7 +155,8 @@ pub(crate) fn remove_argument(
     } else {
         // Case 3: argument or keyword is the last node, so we have to find the last
         // comma in the stmt.
-        for (tok, range) in lexer::lex_starts_at(contents, mode, call_at).flatten() {
+        for (tok, range) in lexer::lex_starts_at(contents, source_type.as_mode(), call_at).flatten()
+        {
             if range.start() == expr_range.start() {
                 fix_end = Some(expr_range.end());
                 break;
diff --git a/crates/ruff/src/checkers/ast/analyze/bindings.rs b/crates/ruff/src/checkers/ast/analyze/bindings.rs
index 80c632308fbea..863da0d9f6dcd 100644
--- a/crates/ruff/src/checkers/ast/analyze/bindings.rs
+++ b/crates/ruff/src/checkers/ast/analyze/bindings.rs
@@ -56,7 +56,7 @@ pub(crate) fn bindings(checker: &mut Checker) {
                 checker.diagnostics.push(diagnostic);
             }
         }
-        if checker.is_stub {
+        if checker.source_type.is_stub() {
             if checker.enabled(Rule::UnaliasedCollectionsAbcSetImport) {
                 if let Some(diagnostic) =
                     flake8_pyi::rules::unaliased_collections_abc_set_import(checker, binding)
diff --git a/crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs b/crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs
index f020fcadbb674..996e433fba25b 100644
--- a/crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs
+++ b/crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs
@@ -37,7 +37,7 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) {
     // Identify any valid runtime imports. If a module is imported at runtime, and
     // used at runtime, then by default, we avoid flagging any other
     // imports from that model as typing-only.
-    let enforce_typing_imports = !checker.is_stub
+    let enforce_typing_imports = !checker.source_type.is_stub()
         && checker.any_enabled(&[
             Rule::RuntimeImportInTypeCheckingBlock,
             Rule::TypingOnlyFirstPartyImport,
@@ -245,7 +245,7 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) {
                 pyflakes::rules::unused_annotation(checker, scope, &mut diagnostics);
             }
 
-            if !checker.is_stub {
+            if !checker.source_type.is_stub() {
                 if checker.any_enabled(&[
                     Rule::UnusedClassMethodArgument,
                     Rule::UnusedFunctionArgument,
diff --git a/crates/ruff/src/checkers/ast/analyze/definitions.rs b/crates/ruff/src/checkers/ast/analyze/definitions.rs
index 83068ca94c739..389695e06ecb6 100644
--- a/crates/ruff/src/checkers/ast/analyze/definitions.rs
+++ b/crates/ruff/src/checkers/ast/analyze/definitions.rs
@@ -30,7 +30,7 @@ pub(crate) fn definitions(checker: &mut Checker) {
         Rule::MissingTypeKwargs,
         Rule::MissingTypeSelf,
     ]);
-    let enforce_stubs = checker.is_stub
+    let enforce_stubs = checker.source_type.is_stub()
         && checker.any_enabled(&[Rule::DocstringInStub, Rule::IterMethodReturnIterable]);
     let enforce_docstrings = checker.any_enabled(&[
         Rule::BlankLineAfterLastSection,
diff --git a/crates/ruff/src/checkers/ast/analyze/expression.rs b/crates/ruff/src/checkers/ast/analyze/expression.rs
index 8199464d7ec03..5519e61284c13 100644
--- a/crates/ruff/src/checkers/ast/analyze/expression.rs
+++ b/crates/ruff/src/checkers/ast/analyze/expression.rs
@@ -31,7 +31,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 if let Some(operator) = typing::to_pep604_operator(value, slice, &checker.semantic)
                 {
                     if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
-                        if !checker.is_stub
+                        if !checker.source_type.is_stub()
                             && checker.settings.target_version < PythonVersion::Py310
                             && checker.settings.target_version >= PythonVersion::Py37
                             && !checker.semantic.future_annotations()
@@ -44,7 +44,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                         }
                     }
                     if checker.enabled(Rule::NonPEP604Annotation) {
-                        if checker.is_stub
+                        if checker.source_type.is_stub()
                             || checker.settings.target_version >= PythonVersion::Py310
                             || (checker.settings.target_version >= PythonVersion::Py37
                                 && checker.semantic.future_annotations()
@@ -59,7 +59,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
 
             // Ex) list[...]
             if checker.enabled(Rule::FutureRequiredTypeAnnotation) {
-                if !checker.is_stub
+                if !checker.source_type.is_stub()
                     && checker.settings.target_version < PythonVersion::Py39
                     && !checker.semantic.future_annotations()
                     && checker.semantic.in_annotation()
@@ -152,7 +152,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                     if checker.enabled(Rule::NumpyDeprecatedFunction) {
                         numpy::rules::deprecated_function(checker, expr);
                     }
-                    if checker.is_stub {
+                    if checker.source_type.is_stub() {
                         if checker.enabled(Rule::CollectionsNamedTuple) {
                             flake8_pyi::rules::collections_named_tuple(checker, expr);
                         }
@@ -167,7 +167,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                             typing::to_pep585_generic(expr, &checker.semantic)
                         {
                             if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
-                                if !checker.is_stub
+                                if !checker.source_type.is_stub()
                                     && checker.settings.target_version < PythonVersion::Py39
                                     && checker.settings.target_version >= PythonVersion::Py37
                                     && !checker.semantic.future_annotations()
@@ -180,7 +180,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                                 }
                             }
                             if checker.enabled(Rule::NonPEP585Annotation) {
-                                if checker.is_stub
+                                if checker.source_type.is_stub()
                                     || checker.settings.target_version >= PythonVersion::Py39
                                     || (checker.settings.target_version >= PythonVersion::Py37
                                         && checker.semantic.future_annotations()
@@ -262,7 +262,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             ]) {
                 if let Some(replacement) = typing::to_pep585_generic(expr, &checker.semantic) {
                     if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
-                        if !checker.is_stub
+                        if !checker.source_type.is_stub()
                             && checker.settings.target_version < PythonVersion::Py39
                             && checker.settings.target_version >= PythonVersion::Py37
                             && !checker.semantic.future_annotations()
@@ -275,7 +275,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                         }
                     }
                     if checker.enabled(Rule::NonPEP585Annotation) {
-                        if checker.is_stub
+                        if checker.source_type.is_stub()
                             || checker.settings.target_version >= PythonVersion::Py39
                             || (checker.settings.target_version >= PythonVersion::Py37
                                 && checker.semantic.future_annotations()
@@ -313,7 +313,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::PrivateMemberAccess) {
                 flake8_self::rules::private_member_access(checker, expr);
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.enabled(Rule::CollectionsNamedTuple) {
                     flake8_pyi::rules::collections_named_tuple(checker, expr);
                 }
@@ -875,7 +875,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::DjangoLocalsInRenderFunction) {
                 flake8_django::rules::locals_in_render_function(checker, func, args, keywords);
             }
-            if checker.is_stub && checker.enabled(Rule::UnsupportedMethodCallOnAll) {
+            if checker.source_type.is_stub() && checker.enabled(Rule::UnsupportedMethodCallOnAll) {
                 flake8_pyi::rules::unsupported_method_call_on_all(checker, func);
             }
         }
@@ -1071,7 +1071,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
         }) => {
             // Ex) `str | None`
             if checker.enabled(Rule::FutureRequiredTypeAnnotation) {
-                if !checker.is_stub
+                if !checker.source_type.is_stub()
                     && checker.settings.target_version < PythonVersion::Py310
                     && !checker.semantic.future_annotations()
                     && checker.semantic.in_annotation()
@@ -1083,7 +1083,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                     );
                 }
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.enabled(Rule::DuplicateUnionMember)
                         && checker.semantic.in_type_definition()
                         // Avoid duplicate checks if the parent is an `|`
@@ -1211,7 +1211,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             kind: _,
             range: _,
         }) => {
-            if checker.is_stub && checker.enabled(Rule::NumericLiteralTooLong) {
+            if checker.source_type.is_stub() && checker.enabled(Rule::NumericLiteralTooLong) {
                 flake8_pyi::rules::numeric_literal_too_long(checker, expr);
             }
         }
@@ -1220,7 +1220,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             kind: _,
             range: _,
         }) => {
-            if checker.is_stub && checker.enabled(Rule::StringOrBytesTooLong) {
+            if checker.source_type.is_stub() && checker.enabled(Rule::StringOrBytesTooLong) {
                 flake8_pyi::rules::string_or_bytes_too_long(checker, expr);
             }
         }
@@ -1248,7 +1248,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UnicodeKindPrefix) {
                 pyupgrade::rules::unicode_kind_prefix(checker, expr, kind.as_deref());
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.enabled(Rule::StringOrBytesTooLong) {
                     flake8_pyi::rules::string_or_bytes_too_long(checker, expr);
                 }
diff --git a/crates/ruff/src/checkers/ast/analyze/parameters.rs b/crates/ruff/src/checkers/ast/analyze/parameters.rs
index 9320947a4539e..fe300a8ea135f 100644
--- a/crates/ruff/src/checkers/ast/analyze/parameters.rs
+++ b/crates/ruff/src/checkers/ast/analyze/parameters.rs
@@ -15,7 +15,7 @@ pub(crate) fn parameters(parameters: &Parameters, checker: &mut Checker) {
     if checker.settings.rules.enabled(Rule::ImplicitOptional) {
         ruff::rules::implicit_optional(checker, parameters);
     }
-    if checker.is_stub {
+    if checker.source_type.is_stub() {
         if checker.enabled(Rule::TypedArgumentDefaultInStub) {
             flake8_pyi::rules::typed_argument_simple_defaults(checker, parameters);
         }
diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs
index 6228e1a4ffb57..e73c79fb08ac6 100644
--- a/crates/ruff/src/checkers/ast/analyze/statement.rs
+++ b/crates/ruff/src/checkers/ast/analyze/statement.rs
@@ -131,7 +131,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     checker.diagnostics.push(diagnostic);
                 }
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.enabled(Rule::PassStatementStubBody) {
                     flake8_pyi::rules::pass_statement_stub_body(checker, body);
                 }
@@ -391,7 +391,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::DjangoUnorderedBodyContentInModel) {
                 flake8_django::rules::unordered_body_content_in_model(checker, bases, body);
             }
-            if !checker.is_stub {
+            if !checker.source_type.is_stub() {
                 if checker.enabled(Rule::DjangoModelWithoutDunderStr) {
                     flake8_django::rules::model_without_dunder_str(checker, class_def);
                 }
@@ -432,7 +432,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     checker.diagnostics.push(diagnostic);
                 }
             }
-            if !checker.is_stub {
+            if !checker.source_type.is_stub() {
                 if checker.any_enabled(&[
                     Rule::AbstractBaseClassWithoutAbstractMethod,
                     Rule::EmptyMethodWithoutAbstractDecorator,
@@ -442,7 +442,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     );
                 }
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.enabled(Rule::PassStatementStubBody) {
                     flake8_pyi::rules::pass_statement_stub_body(checker, body);
                 }
@@ -544,7 +544,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                         alias,
                     );
                 }
-                if !checker.is_stub {
+                if !checker.source_type.is_stub() {
                     if checker.enabled(Rule::UselessImportAlias) {
                         pylint::rules::useless_import_alias(checker, alias);
                     }
@@ -719,7 +719,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     checker.diagnostics.push(diagnostic);
                 }
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.enabled(Rule::FutureAnnotationsInStub) {
                     flake8_pyi::rules::from_future_import(checker, import_from);
                 }
@@ -864,7 +864,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                             checker.diagnostics.push(diagnostic);
                         }
                     }
-                    if !checker.is_stub {
+                    if !checker.source_type.is_stub() {
                         if checker.enabled(Rule::UselessImportAlias) {
                             pylint::rules::useless_import_alias(checker, alias);
                         }
@@ -988,7 +988,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     checker.diagnostics.push(diagnostic);
                 }
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.any_enabled(&[
                     Rule::UnrecognizedVersionInfoCheck,
                     Rule::PatchVersionComparison,
@@ -1300,7 +1300,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.settings.rules.enabled(Rule::TypeBivariance) {
                 pylint::rules::type_bivariance(checker, value);
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.any_enabled(&[
                     Rule::UnprefixedTypeParam,
                     Rule::AssignmentDefaultInStub,
@@ -1365,7 +1365,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     stmt,
                 );
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if let Some(value) = value {
                     if checker.enabled(Rule::AssignmentDefaultInStub) {
                         // Ignore assignments in function bodies; those are covered by other rules.
diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
index e23cc0d700ccd..c7e210076d5b1 100644
--- a/crates/ruff/src/checkers/ast/mod.rs
+++ b/crates/ruff/src/checkers/ast/mod.rs
@@ -52,7 +52,6 @@ use ruff_python_semantic::{
     ModuleKind, ScopeId, ScopeKind, SemanticModel, SemanticModelFlags, StarImport, SubmoduleImport,
 };
 use ruff_python_stdlib::builtins::{BUILTINS, MAGIC_GLOBALS};
-use ruff_python_stdlib::path::{is_jupyter_notebook, is_python_stub_file};
 use ruff_source_file::Locator;
 
 use crate::checkers::ast::deferred::Deferred;
@@ -62,6 +61,7 @@ use crate::noqa::NoqaMapping;
 use crate::registry::Rule;
 use crate::rules::{flake8_pyi, flake8_type_checking, pyflakes, pyupgrade};
 use crate::settings::{flags, Settings};
+use crate::source_kind::PySourceType;
 use crate::{docstrings, noqa};
 
 mod analyze;
@@ -74,10 +74,8 @@ pub(crate) struct Checker<'a> {
     package: Option<&'a Path>,
     /// The module representation of the current file (e.g., `foo.bar`).
     module_path: Option<&'a [String]>,
-    /// Whether the current file is a stub (`.pyi`) file.
-    is_stub: bool,
-    /// Whether the current file is a Jupyter notebook (`.ipynb`) file.
-    pub(crate) is_jupyter_notebook: bool,
+    /// The [`PySourceType`] of the current file.
+    pub(crate) source_type: PySourceType,
     /// The [`flags::Noqa`] for the current analysis (i.e., whether to respect suppression
     /// comments).
     noqa: flags::Noqa,
@@ -119,6 +117,7 @@ impl<'a> Checker<'a> {
         stylist: &'a Stylist,
         indexer: &'a Indexer,
         importer: Importer<'a>,
+        source_type: PySourceType,
     ) -> Checker<'a> {
         Checker {
             settings,
@@ -127,8 +126,7 @@ impl<'a> Checker<'a> {
             path,
             package,
             module_path: module.path(),
-            is_stub: is_python_stub_file(path),
-            is_jupyter_notebook: is_jupyter_notebook(path),
+            source_type,
             locator,
             stylist,
             indexer,
@@ -1788,7 +1786,7 @@ impl<'a> Checker<'a> {
                             pyupgrade::rules::quoted_annotation(self, value, range);
                         }
                     }
-                    if self.is_stub {
+                    if self.source_type.is_stub() {
                         if self.enabled(Rule::QuotedAnnotationInStub) {
                             flake8_pyi::rules::quoted_annotation_in_stub(self, value, range);
                         }
@@ -1930,6 +1928,7 @@ pub(crate) fn check_ast(
     noqa: flags::Noqa,
     path: &Path,
     package: Option<&Path>,
+    source_type: PySourceType,
 ) -> Vec<Diagnostic> {
     let module_path = package.and_then(|package| to_module_path(package, path));
     let module = Module {
@@ -1957,6 +1956,7 @@ pub(crate) fn check_ast(
         stylist,
         indexer,
         Importer::new(python_ast, locator, stylist),
+        source_type,
     );
     checker.bind_builtins();
 
diff --git a/crates/ruff/src/checkers/imports.rs b/crates/ruff/src/checkers/imports.rs
index 7a16036445a51..e263cb55b6754 100644
--- a/crates/ruff/src/checkers/imports.rs
+++ b/crates/ruff/src/checkers/imports.rs
@@ -10,7 +10,7 @@ use ruff_python_ast::imports::{ImportMap, ModuleImport};
 use ruff_python_ast::statement_visitor::StatementVisitor;
 use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
-use ruff_python_stdlib::path::is_python_stub_file;
+
 use ruff_source_file::Locator;
 
 use crate::directives::IsortDirectives;
@@ -18,7 +18,7 @@ use crate::registry::Rule;
 use crate::rules::isort;
 use crate::rules::isort::block::{Block, BlockBuilder};
 use crate::settings::Settings;
-use crate::source_kind::SourceKind;
+use crate::source_kind::{PySourceType, SourceKind};
 
 fn extract_import_map(path: &Path, package: Option<&Path>, blocks: &[&Block]) -> Option<ImportMap> {
     let Some(package) = package else {
@@ -87,12 +87,12 @@ pub(crate) fn check_imports(
     path: &Path,
     package: Option<&Path>,
     source_kind: Option<&SourceKind>,
+    source_type: PySourceType,
 ) -> (Vec<Diagnostic>, Option<ImportMap>) {
-    let is_stub = is_python_stub_file(path);
-
     // Extract all import blocks from the AST.
     let tracker = {
-        let mut tracker = BlockBuilder::new(locator, directives, is_stub, source_kind);
+        let mut tracker =
+            BlockBuilder::new(locator, directives, source_type.is_stub(), source_kind);
         tracker.visit_body(python_ast);
         tracker
     };
@@ -110,7 +110,7 @@ pub(crate) fn check_imports(
                     indexer,
                     settings,
                     package,
-                    source_kind.map_or(false, SourceKind::is_jupyter),
+                    source_type,
                 ) {
                     diagnostics.push(diagnostic);
                 }
@@ -119,7 +119,11 @@ pub(crate) fn check_imports(
     }
     if settings.rules.enabled(Rule::MissingRequiredImport) {
         diagnostics.extend(isort::rules::add_required_imports(
-            python_ast, locator, stylist, settings, is_stub,
+            python_ast,
+            locator,
+            stylist,
+            settings,
+            source_type,
         ));
     }
 
diff --git a/crates/ruff/src/importer/insertion.rs b/crates/ruff/src/importer/insertion.rs
index 3df2781a57ab3..99b58fff065c8 100644
--- a/crates/ruff/src/importer/insertion.rs
+++ b/crates/ruff/src/importer/insertion.rs
@@ -2,7 +2,7 @@
 use std::ops::Add;
 
 use ruff_python_ast::{Ranged, Stmt};
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 use ruff_text_size::TextSize;
 
 use ruff_diagnostics::Edit;
@@ -11,6 +11,8 @@ use ruff_python_codegen::Stylist;
 use ruff_python_trivia::{textwrap::indent, PythonWhitespace};
 use ruff_source_file::{Locator, UniversalNewlineIterator};
 
+use crate::source_kind::PySourceType;
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(super) enum Placement<'a> {
     /// The content will be inserted inline with the existing code (i.e., within semicolon-delimited
@@ -137,7 +139,7 @@ impl<'a> Insertion<'a> {
         mut location: TextSize,
         locator: &Locator<'a>,
         stylist: &Stylist,
-        is_jupyter_notebook: bool,
+        source_type: PySourceType,
     ) -> Insertion<'a> {
         enum Awaiting {
             Colon(u32),
@@ -145,14 +147,9 @@ impl<'a> Insertion<'a> {
             Indent,
         }
 
-        let mode = if is_jupyter_notebook {
-            Mode::Jupyter
-        } else {
-            Mode::Module
-        };
-
         let mut state = Awaiting::Colon(0);
-        for (tok, range) in lexer::lex_starts_at(locator.after(location), mode, location).flatten()
+        for (tok, range) in
+            lexer::lex_starts_at(locator.after(location), source_type.as_mode(), location).flatten()
         {
             match state {
                 // Iterate until we find the colon indicating the start of the block body.
@@ -312,6 +309,8 @@ mod tests {
     use ruff_source_file::{LineEnding, Locator};
     use ruff_text_size::TextSize;
 
+    use crate::source_kind::PySourceType;
+
     use super::Insertion;
 
     #[test]
@@ -432,7 +431,7 @@ x = 1
             let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
-            Insertion::start_of_block(offset, &locator, &stylist, false)
+            Insertion::start_of_block(offset, &locator, &stylist, PySourceType::default())
         }
 
         let contents = "if True: pass";
diff --git a/crates/ruff/src/importer/mod.rs b/crates/ruff/src/importer/mod.rs
index e3860f4fa7ad8..1c317743d3768 100644
--- a/crates/ruff/src/importer/mod.rs
+++ b/crates/ruff/src/importer/mod.rs
@@ -21,6 +21,7 @@ use crate::autofix;
 use crate::autofix::codemods::CodegenStylist;
 use crate::cst::matchers::{match_aliases, match_import_from, match_statement};
 use crate::importer::insertion::Insertion;
+use crate::source_kind::PySourceType;
 
 mod insertion;
 
@@ -121,7 +122,7 @@ impl<'a> Importer<'a> {
         import: &StmtImports,
         at: TextSize,
         semantic: &SemanticModel,
-        is_jupyter_notebook: bool,
+        source_type: PySourceType,
     ) -> Result<TypingImportEdit> {
         // Generate the modified import statement.
         let content = autofix::codemods::retain_imports(
@@ -141,7 +142,7 @@ impl<'a> Importer<'a> {
         // Add the import to a `TYPE_CHECKING` block.
         let add_import_edit = if let Some(block) = self.preceding_type_checking_block(at) {
             // Add the import to the `TYPE_CHECKING` block.
-            self.add_to_type_checking_block(&content, block.start(), is_jupyter_notebook)
+            self.add_to_type_checking_block(&content, block.start(), source_type)
         } else {
             // Add the import to a new `TYPE_CHECKING` block.
             self.add_type_checking_block(
@@ -358,10 +359,9 @@ impl<'a> Importer<'a> {
         &self,
         content: &str,
         at: TextSize,
-        is_jupyter_notebook: bool,
+        source_type: PySourceType,
     ) -> Edit {
-        Insertion::start_of_block(at, self.locator, self.stylist, is_jupyter_notebook)
-            .into_edit(content)
+        Insertion::start_of_block(at, self.locator, self.stylist, source_type).into_edit(content)
     }
 
     /// Return the import statement that precedes the given position, if any.
diff --git a/crates/ruff/src/linter.rs b/crates/ruff/src/linter.rs
index e4a12032e8121..6f9f20bdd40c8 100644
--- a/crates/ruff/src/linter.rs
+++ b/crates/ruff/src/linter.rs
@@ -7,14 +7,14 @@ use colored::Colorize;
 use itertools::Itertools;
 use log::error;
 use ruff_python_parser::lexer::LexResult;
-use ruff_python_parser::{Mode, ParseError};
+use ruff_python_parser::ParseError;
 use rustc_hash::FxHashMap;
 
 use ruff_diagnostics::Diagnostic;
 use ruff_python_ast::imports::ImportMap;
 use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
-use ruff_python_stdlib::path::is_python_stub_file;
+
 use ruff_source_file::{Locator, SourceFileBuilder};
 
 use crate::autofix::{fix_file, FixResult};
@@ -32,7 +32,7 @@ use crate::noqa::add_noqa;
 use crate::registry::{AsRule, Rule};
 use crate::rules::pycodestyle;
 use crate::settings::{flags, Settings};
-use crate::source_kind::SourceKind;
+use crate::source_kind::{PySourceType, SourceKind};
 use crate::{directives, fs};
 
 const CARGO_PKG_NAME: &str = env!("CARGO_PKG_NAME");
@@ -81,6 +81,7 @@ pub fn check_path(
     settings: &Settings,
     noqa: flags::Noqa,
     source_kind: Option<&SourceKind>,
+    source_type: PySourceType,
 ) -> LinterResult<(Vec<Diagnostic>, Option<ImportMap>)> {
     // Aggregate all diagnostics.
     let mut diagnostics = vec![];
@@ -101,9 +102,13 @@ pub fn check_path(
         .iter_enabled()
         .any(|rule_code| rule_code.lint_source().is_tokens())
     {
-        let is_stub = is_python_stub_file(path);
         diagnostics.extend(check_tokens(
-            &tokens, path, locator, indexer, settings, is_stub,
+            &tokens,
+            path,
+            locator,
+            indexer,
+            settings,
+            source_type.is_stub(),
         ));
     }
 
@@ -141,7 +146,7 @@ pub fn check_path(
         match ruff_python_parser::parse_program_tokens(
             tokens,
             &path.to_string_lossy(),
-            source_kind.map_or(false, SourceKind::is_jupyter),
+            source_type.is_jupyter(),
         ) {
             Ok(python_ast) => {
                 if use_ast {
@@ -155,6 +160,7 @@ pub fn check_path(
                         noqa,
                         path,
                         package,
+                        source_type,
                     ));
                 }
                 if use_imports {
@@ -168,6 +174,7 @@ pub fn check_path(
                         path,
                         package,
                         source_kind,
+                        source_type,
                     );
                     imports = module_imports;
                     diagnostics.extend(import_diagnostics);
@@ -260,11 +267,13 @@ const MAX_ITERATIONS: usize = 100;
 
 /// Add any missing `# noqa` pragmas to the source code at the given `Path`.
 pub fn add_noqa_to_path(path: &Path, package: Option<&Path>, settings: &Settings) -> Result<usize> {
+    let source_type = PySourceType::from(path);
+
     // Read the file from disk.
     let contents = std::fs::read_to_string(path)?;
 
     // Tokenize once.
-    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, Mode::Module);
+    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, source_type.as_mode());
 
     // Map row and column locations to byte slices (lazily).
     let locator = Locator::new(&contents);
@@ -298,6 +307,7 @@ pub fn add_noqa_to_path(path: &Path, package: Option<&Path>, settings: &Settings
         settings,
         flags::Noqa::Disabled,
         None,
+        source_type,
     );
 
     // Log any parse errors.
@@ -330,15 +340,10 @@ pub fn lint_only(
     settings: &Settings,
     noqa: flags::Noqa,
     source_kind: Option<&SourceKind>,
+    source_type: PySourceType,
 ) -> LinterResult<(Vec<Message>, Option<ImportMap>)> {
-    let mode = if source_kind.map_or(false, SourceKind::is_jupyter) {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-
     // Tokenize once.
-    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, mode);
+    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, source_type.as_mode());
 
     // Map row and column locations to byte slices (lazily).
     let locator = Locator::new(contents);
@@ -369,6 +374,7 @@ pub fn lint_only(
         settings,
         noqa,
         source_kind,
+        source_type,
     );
 
     result.map(|(diagnostics, imports)| {
@@ -415,6 +421,7 @@ pub fn lint_fix<'a>(
     noqa: flags::Noqa,
     settings: &Settings,
     source_kind: &mut SourceKind,
+    source_type: PySourceType,
 ) -> Result<FixerResult<'a>> {
     let mut transformed = Cow::Borrowed(contents);
 
@@ -427,16 +434,11 @@ pub fn lint_fix<'a>(
     // Track whether the _initial_ source code was parseable.
     let mut parseable = false;
 
-    let mode = if source_kind.is_jupyter() {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-
     // Continuously autofix until the source code stabilizes.
     loop {
         // Tokenize once.
-        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&transformed, mode);
+        let tokens: Vec<LexResult> =
+            ruff_python_parser::tokenize(&transformed, source_type.as_mode());
 
         // Map row and column locations to byte slices (lazily).
         let locator = Locator::new(&transformed);
@@ -467,6 +469,7 @@ pub fn lint_fix<'a>(
             settings,
             noqa,
             Some(source_kind),
+            source_type,
         );
 
         if iterations == 0 {
diff --git a/crates/ruff/src/rules/flake8_annotations/fixes.rs b/crates/ruff/src/rules/flake8_annotations/fixes.rs
index 2481c795126e8..db90d128770a9 100644
--- a/crates/ruff/src/rules/flake8_annotations/fixes.rs
+++ b/crates/ruff/src/rules/flake8_annotations/fixes.rs
@@ -1,29 +1,28 @@
 use anyhow::{bail, Result};
 use ruff_python_ast::{Ranged, Stmt};
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 
 use ruff_diagnostics::Edit;
 use ruff_source_file::Locator;
 
+use crate::source_kind::PySourceType;
+
 /// ANN204
 pub(crate) fn add_return_annotation(
     locator: &Locator,
     stmt: &Stmt,
     annotation: &str,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Result<Edit> {
     let contents = &locator.contents()[stmt.range()];
-    let mode = if is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
 
     // Find the colon (following the `def` keyword).
     let mut seen_lpar = false;
     let mut seen_rpar = false;
     let mut count = 0u32;
-    for (tok, range) in lexer::lex_starts_at(contents, mode, stmt.start()).flatten() {
+    for (tok, range) in
+        lexer::lex_starts_at(contents, source_type.as_mode(), stmt.start()).flatten()
+    {
         if seen_lpar && seen_rpar {
             if matches!(tok, Tok::Colon) {
                 return Ok(Edit::insertion(format!(" -> {annotation}"), range.start()));
diff --git a/crates/ruff/src/rules/flake8_annotations/rules/definition.rs b/crates/ruff/src/rules/flake8_annotations/rules/definition.rs
index 757ab9417ccc2..01895a874dfb7 100644
--- a/crates/ruff/src/rules/flake8_annotations/rules/definition.rs
+++ b/crates/ruff/src/rules/flake8_annotations/rules/definition.rs
@@ -713,7 +713,7 @@ pub(crate) fn definition(
                                 checker.locator(),
                                 stmt,
                                 "None",
-                                checker.is_jupyter_notebook,
+                                checker.source_type,
                             )
                             .map(Fix::suggested)
                         });
@@ -736,7 +736,7 @@ pub(crate) fn definition(
                                 checker.locator(),
                                 stmt,
                                 return_type,
-                                checker.is_jupyter_notebook,
+                                checker.source_type,
                             )
                             .map(Fix::suggested)
                         });
diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs
index e8e8a356c97a2..be620e644b22d 100644
--- a/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs
+++ b/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs
@@ -521,7 +521,7 @@ fn check_fixture_decorator(checker: &mut Checker, func_name: &str, decorator: &D
                                     args,
                                     keywords,
                                     false,
-                                    checker.is_jupyter_notebook,
+                                    checker.source_type,
                                 )
                                 .map(Fix::suggested)
                             });
diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
index 01696bb3ecc85..81c8c329fed04 100644
--- a/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
+++ b/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
@@ -1,5 +1,5 @@
 use ruff_python_ast::{self as ast, Constant, Decorator, Expr, ExprContext, Ranged};
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
@@ -9,6 +9,7 @@ use ruff_source_file::Locator;
 
 use crate::checkers::ast::Checker;
 use crate::registry::{AsRule, Rule};
+use crate::source_kind::PySourceType;
 
 use super::super::types;
 use super::helpers::{is_pytest_parametrize, split_names};
@@ -99,21 +100,19 @@ fn get_parametrize_name_range(
     decorator: &Decorator,
     expr: &Expr,
     locator: &Locator,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> TextRange {
     let mut locations = Vec::new();
     let mut implicit_concat = None;
 
-    let mode = if is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-
     // The parenthesis are not part of the AST, so we need to tokenize the
     // decorator to find them.
-    for (tok, range) in
-        lexer::lex_starts_at(locator.slice(decorator.range()), mode, decorator.start()).flatten()
+    for (tok, range) in lexer::lex_starts_at(
+        locator.slice(decorator.range()),
+        source_type.as_mode(),
+        decorator.start(),
+    )
+    .flatten()
     {
         match tok {
             Tok::Lpar => locations.push(range.start()),
@@ -152,7 +151,7 @@ fn check_names(checker: &mut Checker, decorator: &Decorator, expr: &Expr) {
                             decorator,
                             expr,
                             checker.locator(),
-                            checker.is_jupyter_notebook,
+                            checker.source_type,
                         );
                         let mut diagnostic = Diagnostic::new(
                             PytestParametrizeNamesWrongType {
@@ -187,7 +186,7 @@ fn check_names(checker: &mut Checker, decorator: &Decorator, expr: &Expr) {
                             decorator,
                             expr,
                             checker.locator(),
-                            checker.is_jupyter_notebook,
+                            checker.source_type,
                         );
                         let mut diagnostic = Diagnostic::new(
                             PytestParametrizeNamesWrongType {
diff --git a/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs b/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
index c5b750f793e02..094dc621e8b1a 100644
--- a/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
+++ b/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
@@ -1,5 +1,5 @@
 use ruff_python_ast::{self as ast, Expr, Ranged};
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
@@ -8,6 +8,7 @@ use ruff_source_file::Locator;
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
+use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for unnecessary parentheses on raised exceptions.
@@ -64,7 +65,7 @@ pub(crate) fn unnecessary_paren_on_raise_exception(checker: &mut Checker, expr:
                 return;
             }
 
-            let range = match_parens(func.end(), checker.locator(), checker.is_jupyter_notebook)
+            let range = match_parens(func.end(), checker.locator(), checker.source_type)
                 .expect("Expected call to include parentheses");
             let mut diagnostic = Diagnostic::new(UnnecessaryParenOnRaiseException, range);
             if checker.patch(diagnostic.kind.rule()) {
@@ -79,7 +80,7 @@ pub(crate) fn unnecessary_paren_on_raise_exception(checker: &mut Checker, expr:
 fn match_parens(
     start: TextSize,
     locator: &Locator,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Option<TextRange> {
     let contents = &locator.contents()[usize::from(start)..];
 
@@ -87,13 +88,7 @@ fn match_parens(
     let mut fix_end = None;
     let mut count = 0u32;
 
-    let mode = if is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-
-    for (tok, range) in lexer::lex_starts_at(contents, mode, start).flatten() {
+    for (tok, range) in lexer::lex_starts_at(contents, source_type.as_mode(), start).flatten() {
         match tok {
             Tok::Lpar => {
                 if count == 0 {
diff --git a/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs b/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs
index 9d73345de3633..312c6de97096e 100644
--- a/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs
+++ b/crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs
@@ -374,7 +374,7 @@ pub(crate) fn nested_if_statements(checker: &mut Checker, stmt_if: &StmtIf, pare
     let colon = first_colon_range(
         TextRange::new(test.end(), first_stmt.start()),
         checker.locator().contents(),
-        checker.is_jupyter_notebook,
+        checker.source_type.is_jupyter(),
     );
 
     // Check if the parent is already emitting a larger diagnostic including this if statement
diff --git a/crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs b/crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs
index 26f409ab1d440..55c5d0d78105a 100644
--- a/crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs
+++ b/crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs
@@ -119,7 +119,7 @@ pub(crate) fn multiple_with_statements(
                 body.first().expect("Expected body to be non-empty").start(),
             ),
             checker.locator().contents(),
-            checker.is_jupyter_notebook,
+            checker.source_type.is_jupyter(),
         );
 
         let mut diagnostic = Diagnostic::new(
diff --git a/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs b/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs
index aa7a054001846..5609580ea8113 100644
--- a/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs
+++ b/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs
@@ -447,7 +447,7 @@ fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result
         },
         at,
         checker.semantic(),
-        checker.is_jupyter_notebook,
+        checker.source_type,
     )?;
 
     Ok(
diff --git a/crates/ruff/src/rules/isort/annotate.rs b/crates/ruff/src/rules/isort/annotate.rs
index 351ca6418bc2c..511a152e0b574 100644
--- a/crates/ruff/src/rules/isort/annotate.rs
+++ b/crates/ruff/src/rules/isort/annotate.rs
@@ -3,6 +3,8 @@ use ruff_text_size::TextRange;
 
 use ruff_source_file::Locator;
 
+use crate::source_kind::PySourceType;
+
 use super::comments::Comment;
 use super::helpers::trailing_comma;
 use super::types::{AliasData, TrailingComma};
@@ -13,7 +15,7 @@ pub(crate) fn annotate_imports<'a>(
     comments: Vec<Comment<'a>>,
     locator: &Locator,
     split_on_trailing_comma: bool,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Vec<AnnotatedImport<'a>> {
     let mut comments_iter = comments.into_iter().peekable();
 
@@ -120,7 +122,7 @@ pub(crate) fn annotate_imports<'a>(
                         names: aliases,
                         level: level.map(|level| level.to_u32()),
                         trailing_comma: if split_on_trailing_comma {
-                            trailing_comma(import, locator, is_jupyter_notebook)
+                            trailing_comma(import, locator, source_type)
                         } else {
                             TrailingComma::default()
                         },
diff --git a/crates/ruff/src/rules/isort/comments.rs b/crates/ruff/src/rules/isort/comments.rs
index d0b1564f86497..bd26e5fdb2db1 100644
--- a/crates/ruff/src/rules/isort/comments.rs
+++ b/crates/ruff/src/rules/isort/comments.rs
@@ -1,10 +1,12 @@
 use std::borrow::Cow;
 
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_source_file::Locator;
 
+use crate::source_kind::PySourceType;
+
 #[derive(Debug)]
 pub(crate) struct Comment<'a> {
     pub(crate) value: Cow<'a, str>,
@@ -25,15 +27,10 @@ impl Comment<'_> {
 pub(crate) fn collect_comments<'a>(
     range: TextRange,
     locator: &'a Locator,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Vec<Comment<'a>> {
     let contents = locator.slice(range);
-    let mode = if is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-    lexer::lex_starts_at(contents, mode, range.start())
+    lexer::lex_starts_at(contents, source_type.as_mode(), range.start())
         .flatten()
         .filter_map(|(tok, range)| {
             if let Tok::Comment(value) = tok {
diff --git a/crates/ruff/src/rules/isort/helpers.rs b/crates/ruff/src/rules/isort/helpers.rs
index 89038f7070f28..1e4225c7963ea 100644
--- a/crates/ruff/src/rules/isort/helpers.rs
+++ b/crates/ruff/src/rules/isort/helpers.rs
@@ -1,27 +1,23 @@
 use ruff_python_ast::{Ranged, Stmt};
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 
 use ruff_python_trivia::PythonWhitespace;
 use ruff_source_file::{Locator, UniversalNewlines};
 
 use crate::rules::isort::types::TrailingComma;
+use crate::source_kind::PySourceType;
 
 /// Return `true` if a `Stmt::ImportFrom` statement ends with a magic
 /// trailing comma.
 pub(super) fn trailing_comma(
     stmt: &Stmt,
     locator: &Locator,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> TrailingComma {
     let contents = locator.slice(stmt.range());
     let mut count = 0u32;
     let mut trailing_comma = TrailingComma::Absent;
-    let mode = if is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-    for (tok, _) in lexer::lex_starts_at(contents, mode, stmt.start()).flatten() {
+    for (tok, _) in lexer::lex_starts_at(contents, source_type.as_mode(), stmt.start()).flatten() {
         if matches!(tok, Tok::Lpar) {
             count = count.saturating_add(1);
         }
diff --git a/crates/ruff/src/rules/isort/mod.rs b/crates/ruff/src/rules/isort/mod.rs
index b095b727aeb71..cbb1f55050c21 100644
--- a/crates/ruff/src/rules/isort/mod.rs
+++ b/crates/ruff/src/rules/isort/mod.rs
@@ -22,6 +22,7 @@ use crate::line_width::{LineLength, LineWidth};
 use crate::rules::isort::categorize::KnownModules;
 use crate::rules::isort::types::ImportBlock;
 use crate::settings::types::PythonVersion;
+use crate::source_kind::PySourceType;
 
 mod annotate;
 pub(crate) mod block;
@@ -72,7 +73,7 @@ pub(crate) fn format_imports(
     stylist: &Stylist,
     src: &[PathBuf],
     package: Option<&Path>,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
     combine_as_imports: bool,
     force_single_line: bool,
     force_sort_within_sections: bool,
@@ -100,7 +101,7 @@ pub(crate) fn format_imports(
         comments,
         locator,
         split_on_trailing_comma,
-        is_jupyter_notebook,
+        source_type,
     );
 
     // Normalize imports (i.e., deduplicate, aggregate `from` imports).
diff --git a/crates/ruff/src/rules/isort/rules/add_required_imports.rs b/crates/ruff/src/rules/isort/rules/add_required_imports.rs
index 09ff10669df3a..dbf77d1ba32c2 100644
--- a/crates/ruff/src/rules/isort/rules/add_required_imports.rs
+++ b/crates/ruff/src/rules/isort/rules/add_required_imports.rs
@@ -13,6 +13,7 @@ use ruff_source_file::Locator;
 use crate::importer::Importer;
 use crate::registry::Rule;
 use crate::settings::Settings;
+use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Adds any required imports, as specified by the user, to the top of the
@@ -91,7 +92,7 @@ fn add_required_import(
     locator: &Locator,
     stylist: &Stylist,
     settings: &Settings,
-    is_stub: bool,
+    source_type: PySourceType,
 ) -> Option<Diagnostic> {
     // Don't add imports to semantically-empty files.
     if python_ast.iter().all(is_docstring_stmt) {
@@ -99,7 +100,7 @@ fn add_required_import(
     }
 
     // We don't need to add `__future__` imports to stubs.
-    if is_stub && required_import.is_future_import() {
+    if source_type.is_stub() && required_import.is_future_import() {
         return None;
     }
 
@@ -131,7 +132,7 @@ pub(crate) fn add_required_imports(
     locator: &Locator,
     stylist: &Stylist,
     settings: &Settings,
-    is_stub: bool,
+    source_type: PySourceType,
 ) -> Vec<Diagnostic> {
     settings
         .isort
@@ -172,7 +173,7 @@ pub(crate) fn add_required_imports(
                             locator,
                             stylist,
                             settings,
-                            is_stub,
+                            source_type,
                         )
                     })
                     .collect(),
@@ -190,7 +191,7 @@ pub(crate) fn add_required_imports(
                             locator,
                             stylist,
                             settings,
-                            is_stub,
+                            source_type,
                         )
                     })
                     .collect(),
diff --git a/crates/ruff/src/rules/isort/rules/organize_imports.rs b/crates/ruff/src/rules/isort/rules/organize_imports.rs
index 1cb8182f966e8..175dda8f404c2 100644
--- a/crates/ruff/src/rules/isort/rules/organize_imports.rs
+++ b/crates/ruff/src/rules/isort/rules/organize_imports.rs
@@ -15,6 +15,7 @@ use ruff_source_file::{Locator, UniversalNewlines};
 use crate::line_width::LineWidth;
 use crate::registry::AsRule;
 use crate::settings::Settings;
+use crate::source_kind::PySourceType;
 
 use super::super::block::Block;
 use super::super::{comments, format_imports};
@@ -87,7 +88,7 @@ pub(crate) fn organize_imports(
     indexer: &Indexer,
     settings: &Settings,
     package: Option<&Path>,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Option<Diagnostic> {
     let indentation = locator.slice(extract_indentation_range(&block.imports, locator));
     let indentation = leading_indentation(indentation);
@@ -106,7 +107,7 @@ pub(crate) fn organize_imports(
     let comments = comments::collect_comments(
         TextRange::new(range.start(), locator.full_line_end(range.end())),
         locator,
-        is_jupyter_notebook,
+        source_type,
     );
 
     let trailing_line_end = if block.trailer.is_none() {
@@ -125,7 +126,7 @@ pub(crate) fn organize_imports(
         stylist,
         &settings.src,
         package,
-        is_jupyter_notebook,
+        source_type,
         settings.isort.combine_as_imports,
         settings.isort.force_single_line,
         settings.isort.force_sort_within_sections,
diff --git a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
index 1fb616740e7ef..146d6a4891905 100644
--- a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
+++ b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
@@ -10,6 +10,7 @@ use ruff_source_file::Locator;
 use crate::autofix::edits::remove_argument;
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
+use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for `inplace=True` usages in `pandas` function and method
@@ -107,7 +108,7 @@ pub(crate) fn inplace_argument(
                             keyword.range(),
                             args,
                             keywords,
-                            checker.is_jupyter_notebook,
+                            checker.source_type,
                         ) {
                             diagnostic.set_fix(fix);
                         }
@@ -131,7 +132,7 @@ fn convert_inplace_argument_to_assignment(
     expr_range: TextRange,
     args: &[Expr],
     keywords: &[Keyword],
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Option<Fix> {
     // Add the assignment.
     let call = expr.as_call_expr()?;
@@ -149,7 +150,7 @@ fn convert_inplace_argument_to_assignment(
         args,
         keywords,
         false,
-        is_jupyter_notebook,
+        source_type,
     )
     .ok()?;
 
diff --git a/crates/ruff/src/rules/pyflakes/mod.rs b/crates/ruff/src/rules/pyflakes/mod.rs
index 0d6c7d12eda1b..22cdda5aa4ce4 100644
--- a/crates/ruff/src/rules/pyflakes/mod.rs
+++ b/crates/ruff/src/rules/pyflakes/mod.rs
@@ -12,7 +12,7 @@ mod tests {
     use anyhow::Result;
     use regex::Regex;
     use ruff_python_parser::lexer::LexResult;
-    use ruff_python_parser::Mode;
+    
     use test_case::test_case;
 
     use ruff_diagnostics::Diagnostic;
@@ -25,6 +25,7 @@ mod tests {
     use crate::registry::{AsRule, Linter, Rule};
     use crate::rules::pyflakes;
     use crate::settings::{flags, Settings};
+    use crate::source_kind::PySourceType;
     use crate::test::{test_path, test_snippet};
     use crate::{assert_messages, directives};
 
@@ -505,8 +506,9 @@ mod tests {
     /// Note that all tests marked with `#[ignore]` should be considered TODOs.
     fn flakes(contents: &str, expected: &[Rule]) {
         let contents = dedent(contents);
+        let source_type = PySourceType::default();
         let settings = Settings::for_rules(Linter::Pyflakes.rules());
-        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, Mode::Module);
+        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, source_type.as_mode());
         let locator = Locator::new(&contents);
         let stylist = Stylist::from_tokens(&tokens, &locator);
         let indexer = Indexer::from_tokens(&tokens, &locator);
@@ -530,6 +532,7 @@ mod tests {
             &settings,
             flags::Noqa::Enabled,
             None,
+            source_type,
         );
         diagnostics.sort_by_key(Diagnostic::start);
         let actual = diagnostics
diff --git a/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs b/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
index 9cdec4a422113..97bb87c162c85 100644
--- a/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
+++ b/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
@@ -1,5 +1,5 @@
 use ruff_python_ast::{Expr, Ranged};
-use ruff_python_parser::{lexer, Mode, StringKind, Tok};
+use ruff_python_parser::{lexer, StringKind, Tok};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
@@ -8,6 +8,7 @@ use ruff_source_file::Locator;
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
+use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for f-strings that do not contain any placeholder expressions.
@@ -52,10 +53,10 @@ impl AlwaysAutofixableViolation for FStringMissingPlaceholders {
 fn find_useless_f_strings<'a>(
     expr: &'a Expr,
     locator: &'a Locator,
-    mode: Mode,
+    source_type: PySourceType,
 ) -> impl Iterator<Item = (TextRange, TextRange)> + 'a {
     let contents = locator.slice(expr.range());
-    lexer::lex_starts_at(contents, mode, expr.start())
+    lexer::lex_starts_at(contents, source_type.as_mode(), expr.start())
         .flatten()
         .filter_map(|(tok, range)| match tok {
             Tok::String {
@@ -82,16 +83,13 @@ fn find_useless_f_strings<'a>(
 
 /// F541
 pub(crate) fn f_string_missing_placeholders(expr: &Expr, values: &[Expr], checker: &mut Checker) {
-    let mode = if checker.is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
     if !values
         .iter()
         .any(|value| matches!(value, Expr::FormattedValue(_)))
     {
-        for (prefix_range, tok_range) in find_useless_f_strings(expr, checker.locator(), mode) {
+        for (prefix_range, tok_range) in
+            find_useless_f_strings(expr, checker.locator(), checker.source_type)
+        {
             let mut diagnostic = Diagnostic::new(FStringMissingPlaceholders, tok_range);
             if checker.patch(diagnostic.kind.rule()) {
                 diagnostic.set_fix(convert_f_string_to_regular_string(
diff --git a/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs b/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
index 20544044641b8..15d81589a70a9 100644
--- a/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
+++ b/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
@@ -1,6 +1,6 @@
 use itertools::Itertools;
 use ruff_python_ast::{self as ast, Ranged, Stmt};
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
@@ -12,6 +12,7 @@ use ruff_source_file::Locator;
 use crate::autofix::edits::delete_stmt;
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
+use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for the presence of unused variables in function scopes.
@@ -65,14 +66,14 @@ impl Violation for UnusedVariable {
 fn match_token_before<F>(
     location: TextSize,
     locator: &Locator,
-    mode: Mode,
+    source_type: PySourceType,
     f: F,
 ) -> Option<TextRange>
 where
     F: Fn(Tok) -> bool,
 {
     let contents = locator.after(location);
-    for ((_, range), (tok, _)) in lexer::lex_starts_at(contents, mode, location)
+    for ((_, range), (tok, _)) in lexer::lex_starts_at(contents, source_type.as_mode(), location)
         .flatten()
         .tuple_windows()
     {
@@ -88,7 +89,7 @@ where
 fn match_token_after<F>(
     location: TextSize,
     locator: &Locator,
-    mode: Mode,
+    source_type: PySourceType,
     f: F,
 ) -> Option<TextRange>
 where
@@ -101,7 +102,7 @@ where
     let mut sqb_count = 0u32;
     let mut brace_count = 0u32;
 
-    for ((tok, _), (_, range)) in lexer::lex_starts_at(contents, mode, location)
+    for ((tok, _), (_, range)) in lexer::lex_starts_at(contents, source_type.as_mode(), location)
         .flatten()
         .tuple_windows()
     {
@@ -144,7 +145,7 @@ where
 fn match_token_or_closing_brace<F>(
     location: TextSize,
     locator: &Locator,
-    mode: Mode,
+    source_type: PySourceType,
     f: F,
 ) -> Option<TextRange>
 where
@@ -157,7 +158,7 @@ where
     let mut sqb_count = 0u32;
     let mut brace_count = 0u32;
 
-    for (tok, range) in lexer::lex_starts_at(contents, mode, location).flatten() {
+    for (tok, range) in lexer::lex_starts_at(contents, source_type.as_mode(), location).flatten() {
         match tok {
             Tok::Lpar => {
                 par_count = par_count.saturating_add(1);
@@ -209,12 +210,6 @@ fn remove_unused_variable(
     range: TextRange,
     checker: &Checker,
 ) -> Option<Fix> {
-    let mode = if checker.is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-
     // First case: simple assignment (`x = 1`)
     if let Stmt::Assign(ast::StmtAssign { targets, value, .. }) = stmt {
         if let Some(target) = targets.iter().find(|target| range == target.range()) {
@@ -226,8 +221,10 @@ fn remove_unused_variable(
                     // but preserve the right-hand side.
                     let start = target.start();
                     let end =
-                        match_token_after(start, checker.locator(), mode, |tok| tok == Tok::Equal)?
-                            .start();
+                        match_token_after(start, checker.locator(), checker.source_type, |tok| {
+                            tok == Tok::Equal
+                        })?
+                        .start();
                     let edit = Edit::deletion(start, end);
                     Some(Fix::suggested(edit))
                 } else {
@@ -252,8 +249,10 @@ fn remove_unused_variable(
                 // but preserve the right-hand side.
                 let start = stmt.start();
                 let end =
-                    match_token_after(start, checker.locator(), mode, |tok| tok == Tok::Equal)?
-                        .start();
+                    match_token_after(start, checker.locator(), checker.source_type, |tok| {
+                        tok == Tok::Equal
+                    })?
+                    .start();
                 let edit = Edit::deletion(start, end);
                 Some(Fix::suggested(edit))
             } else {
@@ -275,17 +274,19 @@ fn remove_unused_variable(
                     let start = match_token_before(
                         item.context_expr.start(),
                         checker.locator(),
-                        mode,
+                        checker.source_type,
                         |tok| tok == Tok::As,
                     )?
                     .end();
 
                     // Find the first colon, comma, or closing bracket after the `as` keyword.
-                    let end =
-                        match_token_or_closing_brace(start, checker.locator(), mode, |tok| {
-                            tok == Tok::Colon || tok == Tok::Comma
-                        })?
-                        .start();
+                    let end = match_token_or_closing_brace(
+                        start,
+                        checker.locator(),
+                        checker.source_type,
+                        |tok| tok == Tok::Colon || tok == Tok::Comma,
+                    )?
+                    .start();
 
                     let edit = Edit::deletion(start, end);
                     return Some(Fix::suggested(edit));
diff --git a/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs b/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
index f49cf2aef3f7b..5c5c844bd31cf 100644
--- a/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
+++ b/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
@@ -2,7 +2,7 @@ use std::str::FromStr;
 
 use ruff_python_ast::{self as ast, Constant, Expr, Ranged};
 use ruff_python_literal::cformat::{CFormatPart, CFormatSpec, CFormatStrOrBytes, CFormatString};
-use ruff_python_parser::{lexer, Mode};
+use ruff_python_parser::lexer;
 use ruff_text_size::TextRange;
 use rustc_hash::FxHashMap;
 
@@ -203,12 +203,9 @@ pub(crate) fn bad_string_format_type(checker: &mut Checker, expr: &Expr, right:
     // Grab each string segment (in case there's an implicit concatenation).
     let content = checker.locator().slice(expr.range());
     let mut strings: Vec<TextRange> = vec![];
-    let mode = if checker.is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-    for (tok, range) in lexer::lex_starts_at(content, mode, expr.start()).flatten() {
+    for (tok, range) in
+        lexer::lex_starts_at(content, checker.source_type.as_mode(), expr.start()).flatten()
+    {
         if tok.is_string() {
             strings.push(range);
         } else if tok.is_percent() {
diff --git a/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs b/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
index 6f9fb997dec46..75e208f3182bb 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
@@ -4,7 +4,7 @@ use ruff_python_ast::{self as ast, Constant, Expr, Ranged};
 use ruff_python_literal::cformat::{
     CConversionFlags, CFormatPart, CFormatPrecision, CFormatQuantity, CFormatString,
 };
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
@@ -337,13 +337,12 @@ pub(crate) fn printf_string_formatting(
     // Grab each string segment (in case there's an implicit concatenation).
     let mut strings: Vec<TextRange> = vec![];
     let mut extension = None;
-    let mode = if checker.is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-    for (tok, range) in
-        lexer::lex_starts_at(checker.locator().slice(expr.range()), mode, expr.start()).flatten()
+    for (tok, range) in lexer::lex_starts_at(
+        checker.locator().slice(expr.range()),
+        checker.source_type.as_mode(),
+        expr.start(),
+    )
+    .flatten()
     {
         if tok.is_string() {
             strings.push(range);
diff --git a/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs b/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
index bb7ba78a74367..642bbf910bfec 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
@@ -2,7 +2,7 @@ use std::str::FromStr;
 
 use anyhow::{anyhow, Result};
 use ruff_python_ast::{self as ast, Constant, Expr, Keyword, Ranged};
-use ruff_python_parser::{lexer, Mode};
+use ruff_python_parser::lexer;
 use ruff_text_size::TextSize;
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
@@ -13,6 +13,7 @@ use ruff_source_file::Locator;
 
 use crate::checkers::ast::Checker;
 use crate::registry::Rule;
+use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for redundant `open` mode parameters.
@@ -84,7 +85,7 @@ pub(crate) fn redundant_open_modes(checker: &mut Checker, expr: &Expr) {
                                 mode.replacement_value(),
                                 checker.locator(),
                                 checker.patch(Rule::RedundantOpenModes),
-                                checker.is_jupyter_notebook,
+                                checker.source_type,
                             ));
                         }
                     }
@@ -104,7 +105,7 @@ pub(crate) fn redundant_open_modes(checker: &mut Checker, expr: &Expr) {
                         mode.replacement_value(),
                         checker.locator(),
                         checker.patch(Rule::RedundantOpenModes),
-                        checker.is_jupyter_notebook,
+                        checker.source_type,
                     ));
                 }
             }
@@ -184,7 +185,7 @@ fn create_check(
     replacement_value: Option<&str>,
     locator: &Locator,
     patch: bool,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Diagnostic {
     let mut diagnostic = Diagnostic::new(
         RedundantOpenModes {
@@ -200,8 +201,7 @@ fn create_check(
             )));
         } else {
             diagnostic.try_set_fix(|| {
-                create_remove_param_fix(locator, expr, mode_param, is_jupyter_notebook)
-                    .map(Fix::automatic)
+                create_remove_param_fix(locator, expr, mode_param, source_type).map(Fix::automatic)
             });
         }
     }
@@ -212,7 +212,7 @@ fn create_remove_param_fix(
     locator: &Locator,
     expr: &Expr,
     mode_param: &Expr,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Result<Edit> {
     let content = locator.slice(expr.range());
     // Find the last comma before mode_param and create a deletion fix
@@ -221,12 +221,8 @@ fn create_remove_param_fix(
     let mut fix_end: Option<TextSize> = None;
     let mut is_first_arg: bool = false;
     let mut delete_first_arg: bool = false;
-    let mode = if is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-    for (tok, range) in lexer::lex_starts_at(content, mode, expr.start()).flatten() {
+    for (tok, range) in lexer::lex_starts_at(content, source_type.as_mode(), expr.start()).flatten()
+    {
         if range.start() == mode_param.start() {
             if is_first_arg {
                 delete_first_arg = true;
diff --git a/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs b/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
index 676a390d26506..4b0cae4566d19 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
@@ -9,6 +9,7 @@ use ruff_source_file::Locator;
 use crate::autofix::edits::remove_argument;
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
+use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for uses of `subprocess.run` that send `stdout` and `stderr` to a
@@ -59,7 +60,7 @@ fn generate_fix(
     keywords: &[Keyword],
     stdout: &Keyword,
     stderr: &Keyword,
-    is_jupyter_notebook: bool,
+    source_type: PySourceType,
 ) -> Result<Fix> {
     let (first, second) = if stdout.start() < stderr.start() {
         (stdout, stderr)
@@ -75,7 +76,7 @@ fn generate_fix(
             args,
             keywords,
             false,
-            is_jupyter_notebook,
+            source_type,
         )?],
     ))
 }
@@ -124,7 +125,7 @@ pub(crate) fn replace_stdout_stderr(
                     keywords,
                     stdout,
                     stderr,
-                    checker.is_jupyter_notebook,
+                    checker.source_type,
                 )
             });
         }
diff --git a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
index 658a9fe6b61a2..7f40db0b13463 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
@@ -1,5 +1,5 @@
 use ruff_python_ast::{self as ast, Constant, Expr, Keyword, Ranged};
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_parser::{lexer, Tok};
 use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
@@ -9,6 +9,7 @@ use ruff_source_file::Locator;
 use crate::autofix::edits::remove_argument;
 use crate::checkers::ast::Checker;
 use crate::registry::Rule;
+use crate::source_kind::PySourceType;
 
 #[derive(Debug, PartialEq, Eq)]
 pub(crate) enum Reason {
@@ -122,17 +123,14 @@ fn match_encoding_arg<'a>(args: &'a [Expr], kwargs: &'a [Keyword]) -> Option<Enc
 }
 
 /// Return a [`Fix`] replacing the call to encode with a byte string.
-fn replace_with_bytes_literal(locator: &Locator, expr: &Expr, is_jupyter_notebook: bool) -> Fix {
+fn replace_with_bytes_literal(locator: &Locator, expr: &Expr, source_type: PySourceType) -> Fix {
     // Build up a replacement string by prefixing all string tokens with `b`.
     let contents = locator.slice(expr.range());
     let mut replacement = String::with_capacity(contents.len() + 1);
     let mut prev = expr.start();
-    let mode = if is_jupyter_notebook {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
-    for (tok, range) in lexer::lex_starts_at(contents, mode, expr.start()).flatten() {
+    for (tok, range) in
+        lexer::lex_starts_at(contents, source_type.as_mode(), expr.start()).flatten()
+    {
         match tok {
             Tok::Dot => break,
             Tok::String { .. } => {
@@ -183,7 +181,7 @@ pub(crate) fn unnecessary_encode_utf8(
                         diagnostic.set_fix(replace_with_bytes_literal(
                             checker.locator(),
                             expr,
-                            checker.is_jupyter_notebook,
+                            checker.source_type,
                         ));
                     }
                     checker.diagnostics.push(diagnostic);
@@ -205,7 +203,7 @@ pub(crate) fn unnecessary_encode_utf8(
                                 args,
                                 kwargs,
                                 false,
-                                checker.is_jupyter_notebook,
+                                checker.source_type,
                             )
                             .map(Fix::automatic)
                         });
@@ -228,7 +226,7 @@ pub(crate) fn unnecessary_encode_utf8(
                                 args,
                                 kwargs,
                                 false,
-                                checker.is_jupyter_notebook,
+                                checker.source_type,
                             )
                             .map(Fix::automatic)
                         });
@@ -258,7 +256,7 @@ pub(crate) fn unnecessary_encode_utf8(
                                 args,
                                 kwargs,
                                 false,
-                                checker.is_jupyter_notebook,
+                                checker.source_type,
                             )
                             .map(Fix::automatic)
                         });
@@ -281,7 +279,7 @@ pub(crate) fn unnecessary_encode_utf8(
                                 args,
                                 kwargs,
                                 false,
-                                checker.is_jupyter_notebook,
+                                checker.source_type,
                             )
                             .map(Fix::automatic)
                         });
diff --git a/crates/ruff/src/rules/pyupgrade/rules/useless_object_inheritance.rs b/crates/ruff/src/rules/pyupgrade/rules/useless_object_inheritance.rs
index 82a0e60985362..9e76f846b2541 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/useless_object_inheritance.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/useless_object_inheritance.rs
@@ -73,7 +73,7 @@ pub(crate) fn useless_object_inheritance(checker: &mut Checker, class_def: &ast:
                     &class_def.bases,
                     &class_def.keywords,
                     true,
-                    checker.is_jupyter_notebook,
+                    checker.source_type,
                 )?;
                 Ok(Fix::automatic(edit))
             });
diff --git a/crates/ruff/src/source_kind.rs b/crates/ruff/src/source_kind.rs
index ab63e89c28f2f..3625b32f4e557 100644
--- a/crates/ruff/src/source_kind.rs
+++ b/crates/ruff/src/source_kind.rs
@@ -1,3 +1,7 @@
+use std::path::Path;
+
+use ruff_python_parser::Mode;
+
 use crate::jupyter::Notebook;
 
 #[derive(Clone, Debug, PartialEq, is_macro::Is)]
@@ -24,3 +28,42 @@ impl SourceKind {
         }
     }
 }
+
+#[derive(Clone, Copy, Debug, Default, PartialEq)]
+pub enum PySourceType {
+    #[default]
+    Python,
+    Stub,
+    Jupyter,
+}
+
+impl PySourceType {
+    pub fn as_mode(&self) -> Mode {
+        match self {
+            PySourceType::Python | PySourceType::Stub => Mode::Module,
+            PySourceType::Jupyter => Mode::Jupyter,
+        }
+    }
+
+    pub const fn is_python(&self) -> bool {
+        matches!(self, PySourceType::Python)
+    }
+
+    pub const fn is_stub(&self) -> bool {
+        matches!(self, PySourceType::Stub)
+    }
+
+    pub const fn is_jupyter(&self) -> bool {
+        matches!(self, PySourceType::Jupyter)
+    }
+}
+
+impl From<&Path> for PySourceType {
+    fn from(path: &Path) -> Self {
+        match path.extension() {
+            Some(ext) if ext == "pyi" => PySourceType::Stub,
+            Some(ext) if ext == "ipynb" => PySourceType::Jupyter,
+            _ => PySourceType::Python,
+        }
+    }
+}
diff --git a/crates/ruff/src/test.rs b/crates/ruff/src/test.rs
index f8495064039e6..4fb6f4e9fed64 100644
--- a/crates/ruff/src/test.rs
+++ b/crates/ruff/src/test.rs
@@ -7,7 +7,7 @@ use std::path::Path;
 use anyhow::Result;
 use itertools::Itertools;
 use ruff_python_parser::lexer::LexResult;
-use ruff_python_parser::Mode;
+
 use rustc_hash::FxHashMap;
 
 use ruff_diagnostics::{AutofixKind, Diagnostic};
@@ -26,7 +26,7 @@ use crate::packaging::detect_package_root;
 use crate::registry::AsRule;
 use crate::rules::pycodestyle::rules::syntax_error;
 use crate::settings::{flags, Settings};
-use crate::source_kind::SourceKind;
+use crate::source_kind::{PySourceType, SourceKind};
 
 #[cfg(not(fuzzing))]
 pub(crate) fn read_jupyter_notebook(path: &Path) -> Result<Notebook> {
@@ -101,13 +101,9 @@ pub(crate) fn max_iterations() -> usize {
 /// A convenient wrapper around [`check_path`], that additionally
 /// asserts that autofixes converge after a fixed number of iterations.
 fn test_contents(source_kind: &mut SourceKind, path: &Path, settings: &Settings) -> Vec<Message> {
-    let mode = if source_kind.is_jupyter() {
-        Mode::Jupyter
-    } else {
-        Mode::Module
-    };
     let contents = source_kind.content().to_string();
-    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, mode);
+    let source_type = PySourceType::from(path);
+    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, source_type.as_mode());
     let locator = Locator::new(&contents);
     let stylist = Stylist::from_tokens(&tokens, &locator);
     let indexer = Indexer::from_tokens(&tokens, &locator);
@@ -132,6 +128,7 @@ fn test_contents(source_kind: &mut SourceKind, path: &Path, settings: &Settings)
         settings,
         flags::Noqa::Enabled,
         Some(source_kind),
+        source_type,
     );
 
     let source_has_errors = error.is_some();
@@ -169,7 +166,8 @@ fn test_contents(source_kind: &mut SourceKind, path: &Path, settings: &Settings)
                 notebook.update(&source_map, &fixed_contents);
             };
 
-            let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&fixed_contents, mode);
+            let tokens: Vec<LexResult> =
+                ruff_python_parser::tokenize(&fixed_contents, source_type.as_mode());
             let locator = Locator::new(&fixed_contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             let indexer = Indexer::from_tokens(&tokens, &locator);
@@ -194,6 +192,7 @@ fn test_contents(source_kind: &mut SourceKind, path: &Path, settings: &Settings)
                 settings,
                 flags::Noqa::Enabled,
                 Some(source_kind),
+                source_type,
             );
 
             if let Some(fixed_error) = fixed_error {
diff --git a/crates/ruff_benchmark/benches/linter.rs b/crates/ruff_benchmark/benches/linter.rs
index 7abaa4fdafa25..e7f14ccedbc01 100644
--- a/crates/ruff_benchmark/benches/linter.rs
+++ b/crates/ruff_benchmark/benches/linter.rs
@@ -7,6 +7,7 @@ use criterion::{
 
 use ruff::linter::lint_only;
 use ruff::settings::{flags, Settings};
+use ruff::source_kind::PySourceType;
 use ruff::RuleSelector;
 use ruff_benchmark::{TestCase, TestCaseSpeed, TestFile, TestFileDownloadError};
 
@@ -57,13 +58,15 @@ fn benchmark_linter(mut group: BenchmarkGroup<WallTime>, settings: &Settings) {
             &case,
             |b, case| {
                 b.iter(|| {
+                    let path = case.path();
                     let result = lint_only(
                         case.code(),
-                        &case.path(),
+                        &path,
                         None,
                         settings,
                         flags::Noqa::Enabled,
                         None,
+                        PySourceType::from(path.as_path()),
                     );
 
                     // Assert that file contains no parse errors
diff --git a/crates/ruff_cli/src/diagnostics.rs b/crates/ruff_cli/src/diagnostics.rs
index 54ce93a57dbeb..0553a41a2ae66 100644
--- a/crates/ruff_cli/src/diagnostics.rs
+++ b/crates/ruff_cli/src/diagnostics.rs
@@ -24,12 +24,12 @@ use ruff::message::Message;
 use ruff::pyproject_toml::lint_pyproject_toml;
 use ruff::registry::Rule;
 use ruff::settings::{flags, AllSettings, Settings};
-use ruff::source_kind::SourceKind;
+use ruff::source_kind::{PySourceType, SourceKind};
 use ruff::{fs, IOError};
 use ruff_diagnostics::Diagnostic;
 use ruff_macros::CacheKey;
 use ruff_python_ast::imports::ImportMap;
-use ruff_python_stdlib::path::{is_jupyter_notebook, is_project_toml};
+use ruff_python_stdlib::path::{is_project_toml};
 use ruff_source_file::{LineIndex, SourceCode, SourceFileBuilder};
 
 #[derive(CacheKey)]
@@ -211,8 +211,10 @@ pub(crate) fn lint_path(
         });
     }
 
+    let source_type = PySourceType::from(path);
+
     // Read the file from disk
-    let mut source_kind = if is_jupyter_notebook(path) {
+    let mut source_kind = if source_type.is_jupyter() {
         match load_jupyter_notebook(path) {
             Ok(notebook) => SourceKind::Jupyter(notebook),
             Err(diagnostic) => return Ok(*diagnostic),
@@ -249,6 +251,7 @@ pub(crate) fn lint_path(
             noqa,
             &settings.lib,
             &mut source_kind,
+            source_type,
         ) {
             if !fixed.is_empty() {
                 match autofix {
@@ -335,6 +338,7 @@ pub(crate) fn lint_path(
                 &settings.lib,
                 noqa,
                 Some(&source_kind),
+                source_type,
             );
             let fixed = FxHashMap::default();
             (result, fixed)
@@ -347,6 +351,7 @@ pub(crate) fn lint_path(
             &settings.lib,
             noqa,
             Some(&source_kind),
+            source_type,
         );
         let fixed = FxHashMap::default();
         (result, fixed)
@@ -396,6 +401,8 @@ pub(crate) fn lint_stdin(
     autofix: flags::FixMode,
 ) -> Result<Diagnostics> {
     let mut source_kind = SourceKind::Python(contents.to_string());
+    let source_type = PySourceType::default();
+
     // Lint the inputs.
     let (
         LinterResult {
@@ -415,6 +422,7 @@ pub(crate) fn lint_stdin(
             noqa,
             settings,
             &mut source_kind,
+            source_type,
         ) {
             match autofix {
                 flags::FixMode::Apply => {
@@ -450,6 +458,7 @@ pub(crate) fn lint_stdin(
                 settings,
                 noqa,
                 Some(&source_kind),
+                source_type,
             );
             let fixed = FxHashMap::default();
 
@@ -468,6 +477,7 @@ pub(crate) fn lint_stdin(
             settings,
             noqa,
             Some(&source_kind),
+            source_type,
         );
         let fixed = FxHashMap::default();
         (result, fixed)
diff --git a/crates/ruff_wasm/src/lib.rs b/crates/ruff_wasm/src/lib.rs
index e0e6fc28fecda..85184883612c7 100644
--- a/crates/ruff_wasm/src/lib.rs
+++ b/crates/ruff_wasm/src/lib.rs
@@ -21,6 +21,7 @@ use ruff::rules::{
 use ruff::settings::configuration::Configuration;
 use ruff::settings::options::Options;
 use ruff::settings::{defaults, flags, Settings};
+use ruff::source_kind::PySourceType;
 use ruff_python_codegen::Stylist;
 use ruff_python_formatter::{format_module, format_node, PyFormatOptions};
 use ruff_python_index::{CommentRangesBuilder, Indexer};
@@ -196,8 +197,10 @@ impl Workspace {
     }
 
     pub fn check(&self, contents: &str) -> Result<JsValue, Error> {
+        let source_type = PySourceType::default();
+
         // Tokenize once.
-        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, Mode::Module);
+        let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, source_type.as_mode());
 
         // Map row and column locations to byte slices (lazily).
         let locator = Locator::new(contents);
@@ -227,6 +230,7 @@ impl Workspace {
             &self.settings,
             flags::Noqa::Enabled,
             None,
+            source_type,
         );
 
         let source_code = locator.to_source_code();

From 9d94cd6e92cd5456e37b4440fe7bc1662cb68160 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Fri, 28 Jul 2023 15:14:30 +0530
Subject: [PATCH 11/18] Run cargo fmt

---
 crates/ruff/src/rules/pyflakes/mod.rs | 2 +-
 crates/ruff_cli/src/diagnostics.rs    | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/crates/ruff/src/rules/pyflakes/mod.rs b/crates/ruff/src/rules/pyflakes/mod.rs
index 22cdda5aa4ce4..40b869c930e38 100644
--- a/crates/ruff/src/rules/pyflakes/mod.rs
+++ b/crates/ruff/src/rules/pyflakes/mod.rs
@@ -12,7 +12,7 @@ mod tests {
     use anyhow::Result;
     use regex::Regex;
     use ruff_python_parser::lexer::LexResult;
-    
+
     use test_case::test_case;
 
     use ruff_diagnostics::Diagnostic;
diff --git a/crates/ruff_cli/src/diagnostics.rs b/crates/ruff_cli/src/diagnostics.rs
index 0553a41a2ae66..3120ffe120f9c 100644
--- a/crates/ruff_cli/src/diagnostics.rs
+++ b/crates/ruff_cli/src/diagnostics.rs
@@ -29,7 +29,7 @@ use ruff::{fs, IOError};
 use ruff_diagnostics::Diagnostic;
 use ruff_macros::CacheKey;
 use ruff_python_ast::imports::ImportMap;
-use ruff_python_stdlib::path::{is_project_toml};
+use ruff_python_stdlib::path::is_project_toml;
 use ruff_source_file::{LineIndex, SourceCode, SourceFileBuilder};
 
 #[derive(CacheKey)]

From bc9ec50947961cdb6d9f2238e3dbf479e5f86971 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Fri, 28 Jul 2023 15:15:23 +0530
Subject: [PATCH 12/18] Fix clippy warnings

---
 crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs b/crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs
index 996e433fba25b..e8a5094a83a9d 100644
--- a/crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs
+++ b/crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs
@@ -218,7 +218,7 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) {
             }
         }
 
-        if checker.is_stub {
+        if checker.source_type.is_stub() {
             if checker.enabled(Rule::UnusedPrivateTypeVar) {
                 flake8_pyi::rules::unused_private_type_var(checker, scope, &mut diagnostics);
             }

From 689267bf762aca0a36678d7d2cbf52f41438b046 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 2 Aug 2023 11:32:31 +0530
Subject: [PATCH 13/18] Use `source_type` (leftover from rebase)

---
 crates/ruff/src/checkers/ast/mod.rs                   | 5 -----
 crates/ruff/src/rules/pydocstyle/rules/not_missing.rs | 2 +-
 2 files changed, 1 insertion(+), 6 deletions(-)

diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
index c7e210076d5b1..474e7617c8b64 100644
--- a/crates/ruff/src/checkers/ast/mod.rs
+++ b/crates/ruff/src/checkers/ast/mod.rs
@@ -233,11 +233,6 @@ impl<'a> Checker<'a> {
         &self.semantic
     }
 
-    /// Return `true` if the current file is a stub file (`.pyi`).
-    pub(crate) const fn is_stub(&self) -> bool {
-        self.is_stub
-    }
-
     /// The [`Path`] to the file under analysis.
     pub(crate) const fn path(&self) -> &'a Path {
         self.path
diff --git a/crates/ruff/src/rules/pydocstyle/rules/not_missing.rs b/crates/ruff/src/rules/pydocstyle/rules/not_missing.rs
index d8789d5e0d7b8..00ec74d836efe 100644
--- a/crates/ruff/src/rules/pydocstyle/rules/not_missing.rs
+++ b/crates/ruff/src/rules/pydocstyle/rules/not_missing.rs
@@ -523,7 +523,7 @@ pub(crate) fn not_missing(
     definition: &Definition,
     visibility: Visibility,
 ) -> bool {
-    if checker.is_stub() {
+    if checker.source_type.is_stub() {
         return true;
     }
 

From 65fe5cd3bed5ac3a943d587bc136c9b043999d2b Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Fri, 4 Aug 2023 09:25:55 +0530
Subject: [PATCH 14/18] Fix CI (missed in the merge)

---
 crates/ruff/src/autofix/edits.rs                           | 2 +-
 crates/ruff/src/checkers/ast/analyze/statement.rs          | 4 ++--
 crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/crates/ruff/src/autofix/edits.rs b/crates/ruff/src/autofix/edits.rs
index 8b7761056d9da..da9157140d31b 100644
--- a/crates/ruff/src/autofix/edits.rs
+++ b/crates/ruff/src/autofix/edits.rs
@@ -137,7 +137,7 @@ pub(crate) fn remove_argument<T: Ranged>(
             // previous comma to the end of the argument.
             for (tok, range) in lexer::lex_starts_at(
                 locator.slice(arguments.range()),
-                Mode::Module,
+                source_type.as_mode(),
                 arguments.start(),
             )
             .flatten()
diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs
index ca5181385bb6c..98e52d179ba52 100644
--- a/crates/ruff/src/checkers/ast/analyze/statement.rs
+++ b/crates/ruff/src/checkers/ast/analyze/statement.rs
@@ -168,12 +168,12 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     type_params.as_ref(),
                 );
             }
-            if checker.is_stub {
+            if checker.source_type.is_stub() {
                 if checker.enabled(Rule::StrOrReprDefinedInStub) {
                     flake8_pyi::rules::str_or_repr_defined_in_stub(checker, stmt);
                 }
             }
-            if checker.is_stub || checker.settings.target_version >= PythonVersion::Py311 {
+            if checker.source_type.is_stub() || checker.settings.target_version >= PythonVersion::Py311 {
                 if checker.enabled(Rule::NoReturnArgumentAnnotationInStub) {
                     flake8_pyi::rules::no_return_argument_annotation(checker, parameters);
                 }
diff --git a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
index 4016f354d4c97..d94dc9b955e7d 100644
--- a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
+++ b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
@@ -135,7 +135,7 @@ fn convert_inplace_argument_to_assignment(
         &call.arguments,
         Parentheses::Preserve,
         locator,
-        checker.source_type,
+        source_type,
     )
     .ok()?;
 

From 66edc7f60faadba0e7dd713076ea793843054afb Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Fri, 4 Aug 2023 09:29:30 +0530
Subject: [PATCH 15/18] Run `cargo fmt`

---
 crates/ruff/src/checkers/ast/analyze/statement.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs
index 98e52d179ba52..cbdbff9f5ae54 100644
--- a/crates/ruff/src/checkers/ast/analyze/statement.rs
+++ b/crates/ruff/src/checkers/ast/analyze/statement.rs
@@ -173,7 +173,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     flake8_pyi::rules::str_or_repr_defined_in_stub(checker, stmt);
                 }
             }
-            if checker.source_type.is_stub() || checker.settings.target_version >= PythonVersion::Py311 {
+            if checker.source_type.is_stub()
+                || checker.settings.target_version >= PythonVersion::Py311
+            {
                 if checker.enabled(Rule::NoReturnArgumentAnnotationInStub) {
                     flake8_pyi::rules::no_return_argument_annotation(checker, parameters);
                 }

From ba1fa981022ce09cead26c767213ab47b6fd1ae0 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Sat, 5 Aug 2023 05:50:51 +0530
Subject: [PATCH 16/18] Add `AsMode` trait for `PySourceType`

---
 crates/ruff/src/autofix/edits.rs              |  7 +--
 crates/ruff/src/checkers/ast/mod.rs           |  3 +-
 crates/ruff/src/checkers/imports.rs           |  4 +-
 crates/ruff/src/importer/insertion.rs         |  9 ++--
 crates/ruff/src/importer/mod.rs               |  3 +-
 crates/ruff/src/linter.rs                     |  5 ++-
 .../src/rules/flake8_annotations/fixes.rs     |  6 +--
 .../flake8_pytest_style/rules/parametrize.rs  |  7 +--
 .../unnecessary_paren_on_raise_exception.rs   |  5 +--
 crates/ruff/src/rules/isort/annotate.rs       |  4 +-
 crates/ruff/src/rules/isort/comments.rs       |  5 +--
 crates/ruff/src/rules/isort/helpers.rs        |  5 +--
 crates/ruff/src/rules/isort/mod.rs            |  2 +-
 .../rules/isort/rules/add_required_imports.rs |  3 +-
 .../src/rules/isort/rules/organize_imports.rs |  3 +-
 .../pandas_vet/rules/inplace_argument.rs      |  3 +-
 crates/ruff/src/rules/pyflakes/mod.rs         |  3 +-
 .../rules/f_string_missing_placeholders.rs    |  5 +--
 .../rules/pyflakes/rules/unused_variable.rs   |  5 +--
 .../pylint/rules/bad_string_format_type.rs    |  2 +-
 .../rules/printf_string_formatting.rs         |  2 +-
 .../pyupgrade/rules/redundant_open_modes.rs   |  5 +--
 .../pyupgrade/rules/replace_stdout_stderr.rs  |  3 +-
 .../rules/unnecessary_encode_utf8.rs          |  5 +--
 crates/ruff/src/source_kind.rs                | 43 -------------------
 crates/ruff/src/test.rs                       |  4 +-
 crates/ruff_benchmark/benches/linter.rs       |  2 +-
 crates/ruff_cli/src/diagnostics.rs            |  3 +-
 crates/ruff_python_parser/src/lib.rs          | 15 ++++++-
 crates/ruff_wasm/src/lib.rs                   |  1 +
 30 files changed, 65 insertions(+), 107 deletions(-)

diff --git a/crates/ruff/src/autofix/edits.rs b/crates/ruff/src/autofix/edits.rs
index da9157140d31b..6e90771d8e10b 100644
--- a/crates/ruff/src/autofix/edits.rs
+++ b/crates/ruff/src/autofix/edits.rs
@@ -3,16 +3,17 @@
 use anyhow::{bail, Result};
 
 use ruff_diagnostics::Edit;
-use ruff_python_ast::{self as ast, Arguments, ExceptHandler, Expr, Keyword, Ranged, Stmt};
+use ruff_python_ast::{
+    self as ast, Arguments, ExceptHandler, Expr, Keyword, PySourceType, Ranged, Stmt,
+};
 use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
-use ruff_python_parser::lexer;
+use ruff_python_parser::{lexer, AsMode};
 use ruff_python_trivia::{has_leading_content, is_python_whitespace, PythonWhitespace};
 use ruff_source_file::{Locator, NewlineWithTrailingNewline};
 use ruff_text_size::{TextLen, TextRange, TextSize};
 
 use crate::autofix::codemods;
-use crate::source_kind::PySourceType;
 
 /// Return the `Fix` to use when deleting a `Stmt`.
 ///
diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
index a8243c528bbdf..1c19eda25ea69 100644
--- a/crates/ruff/src/checkers/ast/mod.rs
+++ b/crates/ruff/src/checkers/ast/mod.rs
@@ -43,7 +43,7 @@ use ruff_python_ast::helpers::{extract_handled_exceptions, to_module_path};
 use ruff_python_ast::identifier::Identifier;
 use ruff_python_ast::str::trailing_quote;
 use ruff_python_ast::visitor::{walk_except_handler, walk_pattern, Visitor};
-use ruff_python_ast::{helpers, str, visitor};
+use ruff_python_ast::{helpers, str, visitor, PySourceType};
 use ruff_python_codegen::{Generator, Quote, Stylist};
 use ruff_python_index::Indexer;
 use ruff_python_parser::typing::{parse_type_annotation, AnnotationKind};
@@ -62,7 +62,6 @@ use crate::noqa::NoqaMapping;
 use crate::registry::Rule;
 use crate::rules::{flake8_pyi, flake8_type_checking, pyflakes, pyupgrade};
 use crate::settings::{flags, Settings};
-use crate::source_kind::PySourceType;
 use crate::{docstrings, noqa};
 
 mod analyze;
diff --git a/crates/ruff/src/checkers/imports.rs b/crates/ruff/src/checkers/imports.rs
index e263cb55b6754..d0ee5096577f9 100644
--- a/crates/ruff/src/checkers/imports.rs
+++ b/crates/ruff/src/checkers/imports.rs
@@ -2,7 +2,7 @@
 use std::borrow::Cow;
 use std::path::Path;
 
-use ruff_python_ast::{self as ast, Ranged, Stmt, Suite};
+use ruff_python_ast::{self as ast, PySourceType, Ranged, Stmt, Suite};
 
 use ruff_diagnostics::Diagnostic;
 use ruff_python_ast::helpers::to_module_path;
@@ -18,7 +18,7 @@ use crate::registry::Rule;
 use crate::rules::isort;
 use crate::rules::isort::block::{Block, BlockBuilder};
 use crate::settings::Settings;
-use crate::source_kind::{PySourceType, SourceKind};
+use crate::source_kind::SourceKind;
 
 fn extract_import_map(path: &Path, package: Option<&Path>, blocks: &[&Block]) -> Option<ImportMap> {
     let Some(package) = package else {
diff --git a/crates/ruff/src/importer/insertion.rs b/crates/ruff/src/importer/insertion.rs
index 99b58fff065c8..311a1be96791c 100644
--- a/crates/ruff/src/importer/insertion.rs
+++ b/crates/ruff/src/importer/insertion.rs
@@ -1,8 +1,8 @@
 //! Insert statements into Python code.
 use std::ops::Add;
 
-use ruff_python_ast::{Ranged, Stmt};
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_ast::{PySourceType, Ranged, Stmt};
+use ruff_python_parser::{lexer, AsMode, Tok};
 use ruff_text_size::TextSize;
 
 use ruff_diagnostics::Edit;
@@ -11,8 +11,6 @@ use ruff_python_codegen::Stylist;
 use ruff_python_trivia::{textwrap::indent, PythonWhitespace};
 use ruff_source_file::{Locator, UniversalNewlineIterator};
 
-use crate::source_kind::PySourceType;
-
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(super) enum Placement<'a> {
     /// The content will be inserted inline with the existing code (i.e., within semicolon-delimited
@@ -303,14 +301,13 @@ fn match_leading_semicolon(s: &str) -> Option<TextSize> {
 mod tests {
     use anyhow::Result;
 
+    use ruff_python_ast::PySourceType;
     use ruff_python_codegen::Stylist;
     use ruff_python_parser::lexer::LexResult;
     use ruff_python_parser::{parse_suite, Mode};
     use ruff_source_file::{LineEnding, Locator};
     use ruff_text_size::TextSize;
 
-    use crate::source_kind::PySourceType;
-
     use super::Insertion;
 
     #[test]
diff --git a/crates/ruff/src/importer/mod.rs b/crates/ruff/src/importer/mod.rs
index 1c317743d3768..1ef341ee7ef78 100644
--- a/crates/ruff/src/importer/mod.rs
+++ b/crates/ruff/src/importer/mod.rs
@@ -7,7 +7,7 @@ use std::error::Error;
 
 use anyhow::Result;
 use libcst_native::{ImportAlias, Name, NameOrAttribute};
-use ruff_python_ast::{self as ast, Ranged, Stmt, Suite};
+use ruff_python_ast::{self as ast, PySourceType, Ranged, Stmt, Suite};
 use ruff_text_size::TextSize;
 
 use ruff_diagnostics::Edit;
@@ -21,7 +21,6 @@ use crate::autofix;
 use crate::autofix::codemods::CodegenStylist;
 use crate::cst::matchers::{match_aliases, match_import_from, match_statement};
 use crate::importer::insertion::Insertion;
-use crate::source_kind::PySourceType;
 
 mod insertion;
 
diff --git a/crates/ruff/src/linter.rs b/crates/ruff/src/linter.rs
index 6f9f20bdd40c8..e3fc117253bf2 100644
--- a/crates/ruff/src/linter.rs
+++ b/crates/ruff/src/linter.rs
@@ -7,11 +7,12 @@ use colored::Colorize;
 use itertools::Itertools;
 use log::error;
 use ruff_python_parser::lexer::LexResult;
-use ruff_python_parser::ParseError;
+use ruff_python_parser::{AsMode, ParseError};
 use rustc_hash::FxHashMap;
 
 use ruff_diagnostics::Diagnostic;
 use ruff_python_ast::imports::ImportMap;
+use ruff_python_ast::PySourceType;
 use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
 
@@ -32,7 +33,7 @@ use crate::noqa::add_noqa;
 use crate::registry::{AsRule, Rule};
 use crate::rules::pycodestyle;
 use crate::settings::{flags, Settings};
-use crate::source_kind::{PySourceType, SourceKind};
+use crate::source_kind::SourceKind;
 use crate::{directives, fs};
 
 const CARGO_PKG_NAME: &str = env!("CARGO_PKG_NAME");
diff --git a/crates/ruff/src/rules/flake8_annotations/fixes.rs b/crates/ruff/src/rules/flake8_annotations/fixes.rs
index db90d128770a9..82074eb89a8bd 100644
--- a/crates/ruff/src/rules/flake8_annotations/fixes.rs
+++ b/crates/ruff/src/rules/flake8_annotations/fixes.rs
@@ -1,12 +1,10 @@
 use anyhow::{bail, Result};
-use ruff_python_ast::{Ranged, Stmt};
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_ast::{PySourceType, Ranged, Stmt};
+use ruff_python_parser::{lexer, AsMode, Tok};
 
 use ruff_diagnostics::Edit;
 use ruff_source_file::Locator;
 
-use crate::source_kind::PySourceType;
-
 /// ANN204
 pub(crate) fn add_return_annotation(
     locator: &Locator,
diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
index fd405125f79d9..e42a3d95501a1 100644
--- a/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
+++ b/crates/ruff/src/rules/flake8_pytest_style/rules/parametrize.rs
@@ -1,5 +1,7 @@
-use ruff_python_ast::{self as ast, Arguments, Constant, Decorator, Expr, ExprContext, Ranged};
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_ast::{
+    self as ast, Arguments, Constant, Decorator, Expr, ExprContext, PySourceType, Ranged,
+};
+use ruff_python_parser::{lexer, AsMode, Tok};
 use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
@@ -9,7 +11,6 @@ use ruff_source_file::Locator;
 
 use crate::checkers::ast::Checker;
 use crate::registry::{AsRule, Rule};
-use crate::source_kind::PySourceType;
 
 use super::super::types;
 use super::helpers::{is_pytest_parametrize, split_names};
diff --git a/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs b/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
index a14b69934d0a7..8c0d69bf45d92 100644
--- a/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
+++ b/crates/ruff/src/rules/flake8_raise/rules/unnecessary_paren_on_raise_exception.rs
@@ -1,5 +1,5 @@
-use ruff_python_ast::{self as ast, Arguments, Expr, Ranged};
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_ast::{self as ast, Arguments, Expr, PySourceType, Ranged};
+use ruff_python_parser::{lexer, AsMode, Tok};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
@@ -8,7 +8,6 @@ use ruff_source_file::Locator;
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
-use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for unnecessary parentheses on raised exceptions.
diff --git a/crates/ruff/src/rules/isort/annotate.rs b/crates/ruff/src/rules/isort/annotate.rs
index 511a152e0b574..d434d0bde769f 100644
--- a/crates/ruff/src/rules/isort/annotate.rs
+++ b/crates/ruff/src/rules/isort/annotate.rs
@@ -1,10 +1,8 @@
-use ruff_python_ast::{self as ast, Ranged, Stmt};
+use ruff_python_ast::{self as ast, PySourceType, Ranged, Stmt};
 use ruff_text_size::TextRange;
 
 use ruff_source_file::Locator;
 
-use crate::source_kind::PySourceType;
-
 use super::comments::Comment;
 use super::helpers::trailing_comma;
 use super::types::{AliasData, TrailingComma};
diff --git a/crates/ruff/src/rules/isort/comments.rs b/crates/ruff/src/rules/isort/comments.rs
index bd26e5fdb2db1..e2c05fccc3847 100644
--- a/crates/ruff/src/rules/isort/comments.rs
+++ b/crates/ruff/src/rules/isort/comments.rs
@@ -1,12 +1,11 @@
 use std::borrow::Cow;
 
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_ast::PySourceType;
+use ruff_python_parser::{lexer, AsMode, Tok};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_source_file::Locator;
 
-use crate::source_kind::PySourceType;
-
 #[derive(Debug)]
 pub(crate) struct Comment<'a> {
     pub(crate) value: Cow<'a, str>,
diff --git a/crates/ruff/src/rules/isort/helpers.rs b/crates/ruff/src/rules/isort/helpers.rs
index 1e4225c7963ea..00f9504ac3212 100644
--- a/crates/ruff/src/rules/isort/helpers.rs
+++ b/crates/ruff/src/rules/isort/helpers.rs
@@ -1,11 +1,10 @@
-use ruff_python_ast::{Ranged, Stmt};
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_ast::{PySourceType, Ranged, Stmt};
+use ruff_python_parser::{lexer, AsMode, Tok};
 
 use ruff_python_trivia::PythonWhitespace;
 use ruff_source_file::{Locator, UniversalNewlines};
 
 use crate::rules::isort::types::TrailingComma;
-use crate::source_kind::PySourceType;
 
 /// Return `true` if a `Stmt::ImportFrom` statement ends with a magic
 /// trailing comma.
diff --git a/crates/ruff/src/rules/isort/mod.rs b/crates/ruff/src/rules/isort/mod.rs
index cbb1f55050c21..eed06c0e3d4d5 100644
--- a/crates/ruff/src/rules/isort/mod.rs
+++ b/crates/ruff/src/rules/isort/mod.rs
@@ -11,6 +11,7 @@ pub use categorize::{ImportSection, ImportType};
 use comments::Comment;
 use normalize::normalize_imports;
 use order::order_imports;
+use ruff_python_ast::PySourceType;
 use ruff_python_codegen::Stylist;
 use ruff_source_file::Locator;
 use settings::RelativeImportsOrder;
@@ -22,7 +23,6 @@ use crate::line_width::{LineLength, LineWidth};
 use crate::rules::isort::categorize::KnownModules;
 use crate::rules::isort::types::ImportBlock;
 use crate::settings::types::PythonVersion;
-use crate::source_kind::PySourceType;
 
 mod annotate;
 pub(crate) mod block;
diff --git a/crates/ruff/src/rules/isort/rules/add_required_imports.rs b/crates/ruff/src/rules/isort/rules/add_required_imports.rs
index dbf77d1ba32c2..dd71b1e552638 100644
--- a/crates/ruff/src/rules/isort/rules/add_required_imports.rs
+++ b/crates/ruff/src/rules/isort/rules/add_required_imports.rs
@@ -1,5 +1,5 @@
 use log::error;
-use ruff_python_ast::{self as ast, Stmt, Suite};
+use ruff_python_ast::{self as ast, PySourceType, Stmt, Suite};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Fix};
@@ -13,7 +13,6 @@ use ruff_source_file::Locator;
 use crate::importer::Importer;
 use crate::registry::Rule;
 use crate::settings::Settings;
-use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Adds any required imports, as specified by the user, to the top of the
diff --git a/crates/ruff/src/rules/isort/rules/organize_imports.rs b/crates/ruff/src/rules/isort/rules/organize_imports.rs
index 175dda8f404c2..acc9a0ee31d6c 100644
--- a/crates/ruff/src/rules/isort/rules/organize_imports.rs
+++ b/crates/ruff/src/rules/isort/rules/organize_imports.rs
@@ -1,7 +1,7 @@
 use std::path::Path;
 
 use itertools::{EitherOrBoth, Itertools};
-use ruff_python_ast::{Ranged, Stmt};
+use ruff_python_ast::{PySourceType, Ranged, Stmt};
 use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
@@ -15,7 +15,6 @@ use ruff_source_file::{Locator, UniversalNewlines};
 use crate::line_width::LineWidth;
 use crate::registry::AsRule;
 use crate::settings::Settings;
-use crate::source_kind::PySourceType;
 
 use super::super::block::Block;
 use super::super::{comments, format_imports};
diff --git a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
index d94dc9b955e7d..71cb125d7b94b 100644
--- a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
+++ b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs
@@ -1,14 +1,13 @@
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::is_const_true;
-use ruff_python_ast::{self as ast, Keyword, Ranged};
+use ruff_python_ast::{self as ast, Keyword, PySourceType, Ranged};
 use ruff_python_semantic::{BindingKind, Import};
 use ruff_source_file::Locator;
 
 use crate::autofix::edits::{remove_argument, Parentheses};
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
-use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for `inplace=True` usages in `pandas` function and method
diff --git a/crates/ruff/src/rules/pyflakes/mod.rs b/crates/ruff/src/rules/pyflakes/mod.rs
index 244f8f5138ced..ea7288d5a6921 100644
--- a/crates/ruff/src/rules/pyflakes/mod.rs
+++ b/crates/ruff/src/rules/pyflakes/mod.rs
@@ -16,8 +16,10 @@ mod tests {
     use test_case::test_case;
 
     use ruff_diagnostics::Diagnostic;
+    use ruff_python_ast::PySourceType;
     use ruff_python_codegen::Stylist;
     use ruff_python_index::Indexer;
+    use ruff_python_parser::AsMode;
     use ruff_python_trivia::textwrap::dedent;
     use ruff_source_file::Locator;
 
@@ -25,7 +27,6 @@ mod tests {
     use crate::registry::{AsRule, Linter, Rule};
     use crate::rules::pyflakes;
     use crate::settings::{flags, Settings};
-    use crate::source_kind::PySourceType;
     use crate::test::{test_path, test_snippet};
     use crate::{assert_messages, directives};
 
diff --git a/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs b/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
index 97bb87c162c85..4de5016a3ae76 100644
--- a/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
+++ b/crates/ruff/src/rules/pyflakes/rules/f_string_missing_placeholders.rs
@@ -1,5 +1,5 @@
-use ruff_python_ast::{Expr, Ranged};
-use ruff_python_parser::{lexer, StringKind, Tok};
+use ruff_python_ast::{Expr, PySourceType, Ranged};
+use ruff_python_parser::{lexer, AsMode, StringKind, Tok};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
@@ -8,7 +8,6 @@ use ruff_source_file::Locator;
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
-use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for f-strings that do not contain any placeholder expressions.
diff --git a/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs b/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
index 15d81589a70a9..94cbad3727706 100644
--- a/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
+++ b/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
@@ -1,6 +1,6 @@
 use itertools::Itertools;
-use ruff_python_ast::{self as ast, Ranged, Stmt};
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_ast::{self as ast, PySourceType, Ranged, Stmt};
+use ruff_python_parser::{lexer, AsMode, Tok};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
@@ -12,7 +12,6 @@ use ruff_source_file::Locator;
 use crate::autofix::edits::delete_stmt;
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
-use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for the presence of unused variables in function scopes.
diff --git a/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs b/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
index 5c5c844bd31cf..1f2905751d0fe 100644
--- a/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
+++ b/crates/ruff/src/rules/pylint/rules/bad_string_format_type.rs
@@ -2,7 +2,7 @@ use std::str::FromStr;
 
 use ruff_python_ast::{self as ast, Constant, Expr, Ranged};
 use ruff_python_literal::cformat::{CFormatPart, CFormatSpec, CFormatStrOrBytes, CFormatString};
-use ruff_python_parser::lexer;
+use ruff_python_parser::{lexer, AsMode};
 use ruff_text_size::TextRange;
 use rustc_hash::FxHashMap;
 
diff --git a/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs b/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
index 75e208f3182bb..76192f9c88ab8 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/printf_string_formatting.rs
@@ -4,7 +4,7 @@ use ruff_python_ast::{self as ast, Constant, Expr, Ranged};
 use ruff_python_literal::cformat::{
     CConversionFlags, CFormatPart, CFormatPrecision, CFormatQuantity, CFormatString,
 };
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_parser::{lexer, AsMode, Tok};
 use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
diff --git a/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs b/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
index c2ca07dd82cfb..d09ec7315598d 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/redundant_open_modes.rs
@@ -4,15 +4,14 @@ use anyhow::{anyhow, Result};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::{self as ast, Constant, Expr, Ranged};
-use ruff_python_parser::lexer;
+use ruff_python_ast::{self as ast, Constant, Expr, PySourceType, Ranged};
+use ruff_python_parser::{lexer, AsMode};
 use ruff_python_semantic::SemanticModel;
 use ruff_source_file::Locator;
 use ruff_text_size::TextSize;
 
 use crate::checkers::ast::Checker;
 use crate::registry::Rule;
-use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for redundant `open` mode parameters.
diff --git a/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs b/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
index 18e6f9bd1586a..3482e99006488 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs
@@ -2,13 +2,12 @@ use anyhow::Result;
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::{self as ast, Keyword, Ranged};
+use ruff_python_ast::{self as ast, Keyword, PySourceType, Ranged};
 use ruff_source_file::Locator;
 
 use crate::autofix::edits::{remove_argument, Parentheses};
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
-use crate::source_kind::PySourceType;
 
 /// ## What it does
 /// Checks for uses of `subprocess.run` that send `stdout` and `stderr` to a
diff --git a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
index 37313e69adae2..8e9dae557e376 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_encode_utf8.rs
@@ -1,14 +1,13 @@
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::{self as ast, Arguments, Constant, Expr, Keyword, Ranged};
-use ruff_python_parser::{lexer, Tok};
+use ruff_python_ast::{self as ast, Arguments, Constant, Expr, Keyword, PySourceType, Ranged};
+use ruff_python_parser::{lexer, AsMode, Tok};
 use ruff_source_file::Locator;
 use ruff_text_size::TextRange;
 
 use crate::autofix::edits::{remove_argument, Parentheses};
 use crate::checkers::ast::Checker;
 use crate::registry::Rule;
-use crate::source_kind::PySourceType;
 
 #[derive(Debug, PartialEq, Eq)]
 pub(crate) enum Reason {
diff --git a/crates/ruff/src/source_kind.rs b/crates/ruff/src/source_kind.rs
index 3625b32f4e557..ab63e89c28f2f 100644
--- a/crates/ruff/src/source_kind.rs
+++ b/crates/ruff/src/source_kind.rs
@@ -1,7 +1,3 @@
-use std::path::Path;
-
-use ruff_python_parser::Mode;
-
 use crate::jupyter::Notebook;
 
 #[derive(Clone, Debug, PartialEq, is_macro::Is)]
@@ -28,42 +24,3 @@ impl SourceKind {
         }
     }
 }
-
-#[derive(Clone, Copy, Debug, Default, PartialEq)]
-pub enum PySourceType {
-    #[default]
-    Python,
-    Stub,
-    Jupyter,
-}
-
-impl PySourceType {
-    pub fn as_mode(&self) -> Mode {
-        match self {
-            PySourceType::Python | PySourceType::Stub => Mode::Module,
-            PySourceType::Jupyter => Mode::Jupyter,
-        }
-    }
-
-    pub const fn is_python(&self) -> bool {
-        matches!(self, PySourceType::Python)
-    }
-
-    pub const fn is_stub(&self) -> bool {
-        matches!(self, PySourceType::Stub)
-    }
-
-    pub const fn is_jupyter(&self) -> bool {
-        matches!(self, PySourceType::Jupyter)
-    }
-}
-
-impl From<&Path> for PySourceType {
-    fn from(path: &Path) -> Self {
-        match path.extension() {
-            Some(ext) if ext == "pyi" => PySourceType::Stub,
-            Some(ext) if ext == "ipynb" => PySourceType::Jupyter,
-            _ => PySourceType::Python,
-        }
-    }
-}
diff --git a/crates/ruff/src/test.rs b/crates/ruff/src/test.rs
index 4fb6f4e9fed64..b89ce5a80fd93 100644
--- a/crates/ruff/src/test.rs
+++ b/crates/ruff/src/test.rs
@@ -11,8 +11,10 @@ use ruff_python_parser::lexer::LexResult;
 use rustc_hash::FxHashMap;
 
 use ruff_diagnostics::{AutofixKind, Diagnostic};
+use ruff_python_ast::PySourceType;
 use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
+use ruff_python_parser::AsMode;
 use ruff_python_trivia::textwrap::dedent;
 use ruff_source_file::{Locator, SourceFileBuilder};
 
@@ -26,7 +28,7 @@ use crate::packaging::detect_package_root;
 use crate::registry::AsRule;
 use crate::rules::pycodestyle::rules::syntax_error;
 use crate::settings::{flags, Settings};
-use crate::source_kind::{PySourceType, SourceKind};
+use crate::source_kind::SourceKind;
 
 #[cfg(not(fuzzing))]
 pub(crate) fn read_jupyter_notebook(path: &Path) -> Result<Notebook> {
diff --git a/crates/ruff_benchmark/benches/linter.rs b/crates/ruff_benchmark/benches/linter.rs
index e7f14ccedbc01..da0875640b7dd 100644
--- a/crates/ruff_benchmark/benches/linter.rs
+++ b/crates/ruff_benchmark/benches/linter.rs
@@ -7,8 +7,8 @@ use criterion::{
 
 use ruff::linter::lint_only;
 use ruff::settings::{flags, Settings};
-use ruff::source_kind::PySourceType;
 use ruff::RuleSelector;
+use ruff_python_ast::PySourceType;
 use ruff_benchmark::{TestCase, TestCaseSpeed, TestFile, TestFileDownloadError};
 
 #[cfg(target_os = "windows")]
diff --git a/crates/ruff_cli/src/diagnostics.rs b/crates/ruff_cli/src/diagnostics.rs
index 3120ffe120f9c..5ad6a50fd2788 100644
--- a/crates/ruff_cli/src/diagnostics.rs
+++ b/crates/ruff_cli/src/diagnostics.rs
@@ -24,11 +24,12 @@ use ruff::message::Message;
 use ruff::pyproject_toml::lint_pyproject_toml;
 use ruff::registry::Rule;
 use ruff::settings::{flags, AllSettings, Settings};
-use ruff::source_kind::{PySourceType, SourceKind};
+use ruff::source_kind::SourceKind;
 use ruff::{fs, IOError};
 use ruff_diagnostics::Diagnostic;
 use ruff_macros::CacheKey;
 use ruff_python_ast::imports::ImportMap;
+use ruff_python_ast::PySourceType;
 use ruff_python_stdlib::path::is_project_toml;
 use ruff_source_file::{LineIndex, SourceCode, SourceFileBuilder};
 
diff --git a/crates/ruff_python_parser/src/lib.rs b/crates/ruff_python_parser/src/lib.rs
index c30fa66d79236..0e0d607a12768 100644
--- a/crates/ruff_python_parser/src/lib.rs
+++ b/crates/ruff_python_parser/src/lib.rs
@@ -114,7 +114,7 @@ pub use parser::{
     parse, parse_expression, parse_expression_starts_at, parse_program, parse_starts_at,
     parse_suite, parse_tokens, ParseError, ParseErrorType,
 };
-use ruff_python_ast::{CmpOp, Expr, Mod, Ranged, Suite};
+use ruff_python_ast::{CmpOp, Expr, Mod, PySourceType, Ranged, Suite};
 use ruff_text_size::{TextRange, TextSize};
 pub use string::FStringErrorType;
 pub use token::{StringKind, Tok, TokenKind};
@@ -323,6 +323,19 @@ impl std::str::FromStr for Mode {
     }
 }
 
+pub trait AsMode {
+    fn as_mode(&self) -> Mode;
+}
+
+impl AsMode for PySourceType {
+    fn as_mode(&self) -> Mode {
+        match self {
+            PySourceType::Python | PySourceType::Stub => Mode::Module,
+            PySourceType::Jupyter => Mode::Jupyter,
+        }
+    }
+}
+
 /// Returned when a given mode is not valid.
 #[derive(Debug)]
 pub struct ModeParseError;
diff --git a/crates/ruff_wasm/src/lib.rs b/crates/ruff_wasm/src/lib.rs
index 82be35ea07c49..9497a5449b4ea 100644
--- a/crates/ruff_wasm/src/lib.rs
+++ b/crates/ruff_wasm/src/lib.rs
@@ -22,6 +22,7 @@ use ruff::settings::configuration::Configuration;
 use ruff::settings::options::Options;
 use ruff::settings::{defaults, flags, Settings};
 use ruff_python_ast::PySourceType;
+use ruff_python_parser::AsMode;
 use ruff_python_codegen::Stylist;
 use ruff_python_formatter::{format_module, format_node, PyFormatOptions};
 use ruff_python_index::{CommentRangesBuilder, Indexer};

From 851767e37ecba80ae5f35a6d198d4f506792db04 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Sat, 5 Aug 2023 05:52:53 +0530
Subject: [PATCH 17/18] Run `cargo fmt`

---
 crates/ruff_benchmark/benches/linter.rs | 2 +-
 crates/ruff_wasm/src/lib.rs             | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/crates/ruff_benchmark/benches/linter.rs b/crates/ruff_benchmark/benches/linter.rs
index da0875640b7dd..a93f5dbb2817a 100644
--- a/crates/ruff_benchmark/benches/linter.rs
+++ b/crates/ruff_benchmark/benches/linter.rs
@@ -8,8 +8,8 @@ use criterion::{
 use ruff::linter::lint_only;
 use ruff::settings::{flags, Settings};
 use ruff::RuleSelector;
-use ruff_python_ast::PySourceType;
 use ruff_benchmark::{TestCase, TestCaseSpeed, TestFile, TestFileDownloadError};
+use ruff_python_ast::PySourceType;
 
 #[cfg(target_os = "windows")]
 #[global_allocator]
diff --git a/crates/ruff_wasm/src/lib.rs b/crates/ruff_wasm/src/lib.rs
index 9497a5449b4ea..d8aa2c0e64193 100644
--- a/crates/ruff_wasm/src/lib.rs
+++ b/crates/ruff_wasm/src/lib.rs
@@ -22,10 +22,10 @@ use ruff::settings::configuration::Configuration;
 use ruff::settings::options::Options;
 use ruff::settings::{defaults, flags, Settings};
 use ruff_python_ast::PySourceType;
-use ruff_python_parser::AsMode;
 use ruff_python_codegen::Stylist;
 use ruff_python_formatter::{format_module, format_node, PyFormatOptions};
 use ruff_python_index::{CommentRangesBuilder, Indexer};
+use ruff_python_parser::AsMode;
 use ruff_source_file::{Locator, SourceLocation};
 
 #[wasm_bindgen(typescript_custom_section)]

From fd2874ad1cc45e5444251664b25cccb9a672b145 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Thu, 27 Jul 2023 10:42:39 +0530
Subject: [PATCH 18/18] Update `F841` autofix to not remove line magic expr

---
 .../fixtures/jupyter/unused_variable.ipynb    | 49 +++++++++++++
 .../jupyter/unused_variable_expected.ipynb    | 49 +++++++++++++
 crates/ruff/src/jupyter/notebook.rs           | 12 ++++
 ...ter__notebook__tests__unused_variable.snap | 72 +++++++++++++++++++
 crates/ruff_python_ast/src/helpers.rs         |  1 +
 5 files changed, 183 insertions(+)
 create mode 100644 crates/ruff/resources/test/fixtures/jupyter/unused_variable.ipynb
 create mode 100644 crates/ruff/resources/test/fixtures/jupyter/unused_variable_expected.ipynb
 create mode 100644 crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__unused_variable.snap

diff --git a/crates/ruff/resources/test/fixtures/jupyter/unused_variable.ipynb b/crates/ruff/resources/test/fixtures/jupyter/unused_variable.ipynb
new file mode 100644
index 0000000000000..d4b2f8f8af78d
--- /dev/null
+++ b/crates/ruff/resources/test/fixtures/jupyter/unused_variable.ipynb
@@ -0,0 +1,49 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a0efffbc-85f1-4513-bf49-5387ec3a2a4e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def f():\n",
+    "    foo1 = %matplotlib --list\n",
+    "    foo2: list[str] = %matplotlib --list"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6e0b2b50-43f2-4f59-951d-9404dd560ae4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def f():\n",
+    "    bar1 = !pwd\n",
+    "    bar2: str = !pwd"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (ruff)",
+   "language": "python",
+   "name": "ruff"
+  },
+  "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.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/crates/ruff/resources/test/fixtures/jupyter/unused_variable_expected.ipynb b/crates/ruff/resources/test/fixtures/jupyter/unused_variable_expected.ipynb
new file mode 100644
index 0000000000000..0c68e77de6179
--- /dev/null
+++ b/crates/ruff/resources/test/fixtures/jupyter/unused_variable_expected.ipynb
@@ -0,0 +1,49 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24426ef2-046c-453e-b809-05b56e7355e0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def f():\n",
+    "    %matplotlib --list\n",
+    "    %matplotlib --list"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3d98fdae-b86b-476e-b4db-9d3ce5562682",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def f():\n",
+    "    !pwd\n",
+    "    !pwd"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (ruff)",
+   "language": "python",
+   "name": "ruff"
+  },
+  "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.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/crates/ruff/src/jupyter/notebook.rs b/crates/ruff/src/jupyter/notebook.rs
index 652e54850f05f..4aaff0b1ff734 100644
--- a/crates/ruff/src/jupyter/notebook.rs
+++ b/crates/ruff/src/jupyter/notebook.rs
@@ -582,6 +582,18 @@ print("after empty cells")
         Ok(())
     }
 
+    #[test]
+    fn test_unused_variable() -> Result<()> {
+        let path = "unused_variable.ipynb".to_string();
+        let (diagnostics, source_kind, _) = test_notebook_path(
+            &path,
+            Path::new("unused_variable_expected.ipynb"),
+            &settings::Settings::for_rule(Rule::UnusedVariable),
+        )?;
+        assert_messages!(diagnostics, path, source_kind);
+        Ok(())
+    }
+
     #[test]
     fn test_json_consistency() -> Result<()> {
         let path = "before_fix.ipynb".to_string();
diff --git a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__unused_variable.snap b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__unused_variable.snap
new file mode 100644
index 0000000000000..689e20e25ee18
--- /dev/null
+++ b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__unused_variable.snap
@@ -0,0 +1,72 @@
+---
+source: crates/ruff/src/jupyter/notebook.rs
+---
+unused_variable.ipynb:cell 1:2:5: F841 [*] Local variable `foo1` is assigned to but never used
+  |
+1 | def f():
+2 |     foo1 = %matplotlib --list
+  |     ^^^^ F841
+3 |     foo2: list[str] = %matplotlib --list
+  |
+  = help: Remove assignment to unused variable `foo1`
+
+ℹ Suggested fix
+1 1 | def f():
+2   |-    foo1 = %matplotlib --list
+  2 |+    %matplotlib --list
+3 3 |     foo2: list[str] = %matplotlib --list
+4 4 | def f():
+5 5 |     bar1 = !pwd
+
+unused_variable.ipynb:cell 1:3:5: F841 [*] Local variable `foo2` is assigned to but never used
+  |
+1 | def f():
+2 |     foo1 = %matplotlib --list
+3 |     foo2: list[str] = %matplotlib --list
+  |     ^^^^ F841
+  |
+  = help: Remove assignment to unused variable `foo2`
+
+ℹ Suggested fix
+1 1 | def f():
+2 2 |     foo1 = %matplotlib --list
+3   |-    foo2: list[str] = %matplotlib --list
+  3 |+    %matplotlib --list
+4 4 | def f():
+5 5 |     bar1 = !pwd
+6 6 |     bar2: str = !pwd
+
+unused_variable.ipynb:cell 2:2:5: F841 [*] Local variable `bar1` is assigned to but never used
+  |
+1 | def f():
+2 |     bar1 = !pwd
+  |     ^^^^ F841
+3 |     bar2: str = !pwd
+  |
+  = help: Remove assignment to unused variable `bar1`
+
+ℹ Suggested fix
+2 2 |     foo1 = %matplotlib --list
+3 3 |     foo2: list[str] = %matplotlib --list
+4 4 | def f():
+5   |-    bar1 = !pwd
+  5 |+    !pwd
+6 6 |     bar2: str = !pwd
+
+unused_variable.ipynb:cell 2:3:5: F841 [*] Local variable `bar2` is assigned to but never used
+  |
+1 | def f():
+2 |     bar1 = !pwd
+3 |     bar2: str = !pwd
+  |     ^^^^ F841
+  |
+  = help: Remove assignment to unused variable `bar2`
+
+ℹ Suggested fix
+3 3 |     foo2: list[str] = %matplotlib --list
+4 4 | def f():
+5 5 |     bar1 = !pwd
+6   |-    bar2: str = !pwd
+  6 |+    !pwd
+
+
diff --git a/crates/ruff_python_ast/src/helpers.rs b/crates/ruff_python_ast/src/helpers.rs
index a445d624e01b6..2666f32de3452 100644
--- a/crates/ruff_python_ast/src/helpers.rs
+++ b/crates/ruff_python_ast/src/helpers.rs
@@ -111,6 +111,7 @@ where
                 | Expr::Subscript(_)
                 | Expr::Yield(_)
                 | Expr::YieldFrom(_)
+                | Expr::LineMagic(_)
         )
     })
 }