Use Jupyter mode while parsing Notebook files (#5552)

## Summary Enable using the new `Mode::Jupyter` for the tokenizer/parser to parse Jupyter line magic tokens. The individual call to the lexer i.e., `lex_starts_at` done by various rules should consider the context of the source code (is this content from a Jupyter Notebook?). Thus, a new field `source_type` (of type `PySourceType`) is added to `Checker` which is being passed around as an argument to the relevant functions. This is then used to determine the `Mode` for the lexer. ## Test Plan Add new test cases to make sure that the magic statement is considered while generating the diagnostic and autofix: * For `I001`, if there's a magic statement in between two import blocks, they should be sorted independently fixes: #6090
astral-sh · Aug 5, 2023 · 32fa057 · 32fa057
1 parent d788957
commit 32fa057
Show file tree

Hide file tree

Showing 52 changed files with 652 additions and 196 deletions.
diff --git a/crates/ruff/resources/test/fixtures/jupyter/cell/cell_magic.json 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/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
@@ -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/resources/test/fixtures/jupyter/line_magics.ipynb b/crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb
@@ -0,0 +1,52 @@
+{
+ "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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2b0e2986-1b87-4bb6-9b1d-c11ca1decd87",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%timeit\n",
+    "import sys"
+   ]
+  }
+ ],
+ "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
@@ -0,0 +1,51 @@
+{
+ "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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d7b8e967-8b4a-493b-b6f7-d5cecfb3a5c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%timeit\n",
+    "import sys"
+   ]
+  }
+ ],
+ "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/autofix/edits.rs b/crates/ruff/src/autofix/edits.rs
@@ -3,10 +3,12 @@
 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, Mode};
+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};
@@ -88,6 +90,7 @@ pub(crate) fn remove_argument<T: Ranged>(
     arguments: &Arguments,
     parentheses: Parentheses,
     locator: &Locator,
+    source_type: PySourceType,
 ) -> Result<Edit> {
     // TODO(sbrugman): Preserve trailing comments.
     if arguments.keywords.len() + arguments.args.len() > 1 {
@@ -106,7 +109,7 @@ pub(crate) fn remove_argument<T: Ranged>(
             let mut seen_comma = false;
             for (tok, range) in lexer::lex_starts_at(
                 locator.slice(arguments.range()),
-                Mode::Module,
+                source_type.as_mode(),
                 arguments.start(),
             )
             .flatten()
@@ -135,7 +138,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/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,
@@ -243,7 +243,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
@@ -30,7 +30,7 @@ pub(crate) fn definitions(checker: &mut Checker) {
         Rule::MissingTypeKwargs,
         Rule::MissingTypeSelf,
     ]);
-    let enforce_stubs = checker.is_stub && checker.enabled(Rule::DocstringInStub);
+    let enforce_stubs = checker.source_type.is_stub() && checker.enabled(Rule::DocstringInStub);
     let enforce_stubs_and_runtime = checker.enabled(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
@@ -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()
@@ -176,7 +176,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()
@@ -187,7 +187,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()
@@ -272,7 +272,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()
@@ -285,7 +285,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()
@@ -1066,7 +1066,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()
@@ -1212,7 +1212,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);
             }
         }
@@ -1221,7 +1221,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);
             }
         }
@@ -1249,7 +1249,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
@@ -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);
         }