Skip to content

Commit 32fa057

Browse files
dhruvmaniladhruvmanila
authored
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
1 parent d788957 commit 32fa057

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

52 files changed

+652
-196
lines changed

‎crates/ruff/resources/test/fixtures/jupyter/cell/cell_magic.json

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
{
2+
"execution_count": null,
3+
"cell_type": "code",
4+
"id": "1",
5+
"metadata": {},
6+
"outputs": [],
7+
"source": ["%%timeit\n", "print('hello world')"]
8+
}

‎crates/ruff/resources/test/fixtures/jupyter/isort.ipynb

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,23 @@
2525
"def foo():\n",
2626
" pass"
2727
]
28+
},
29+
{
30+
"cell_type": "code",
31+
"execution_count": null,
32+
"id": "16214f6f-bb32-4594-81be-79fb27c6ec92",
33+
"metadata": {},
34+
"outputs": [],
35+
"source": [
36+
"from pathlib import Path\n",
37+
"import sys\n",
38+
"\n",
39+
"%matplotlib \\\n",
40+
" --inline\n",
41+
"\n",
42+
"import math\n",
43+
"import abc"
44+
]
2845
}
2946
],
3047
"metadata": {

‎crates/ruff/resources/test/fixtures/jupyter/isort_expected.ipynb

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,23 @@
2727
"def foo():\n",
2828
" pass"
2929
]
30+
},
31+
{
32+
"cell_type": "code",
33+
"execution_count": null,
34+
"id": "6d6c55c6-4a34-4662-914b-4ee11c9c24a5",
35+
"metadata": {},
36+
"outputs": [],
37+
"source": [
38+
"import sys\n",
39+
"from pathlib import Path\n",
40+
"\n",
41+
"%matplotlib \\\n",
42+
" --inline\n",
43+
"\n",
44+
"import abc\n",
45+
"import math"
46+
]
3047
}
3148
],
3249
"metadata": {

‎crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"id": "eab4754a-d6df-4b41-8ee8-7e23aef440f9",
7+
"metadata": {},
8+
"outputs": [],
9+
"source": [
10+
"import math\n",
11+
"\n",
12+
"%matplotlib inline\n",
13+
"\n",
14+
"import os\n",
15+
"\n",
16+
"_ = math.pi"
17+
]
18+
},
19+
{
20+
"cell_type": "code",
21+
"execution_count": null,
22+
"id": "2b0e2986-1b87-4bb6-9b1d-c11ca1decd87",
23+
"metadata": {},
24+
"outputs": [],
25+
"source": [
26+
"%%timeit\n",
27+
"import sys"
28+
]
29+
}
30+
],
31+
"metadata": {
32+
"kernelspec": {
33+
"display_name": "Python (ruff)",
34+
"language": "python",
35+
"name": "ruff"
36+
},
37+
"language_info": {
38+
"codemirror_mode": {
39+
"name": "ipython",
40+
"version": 3
41+
},
42+
"file_extension": ".py",
43+
"mimetype": "text/x-python",
44+
"name": "python",
45+
"nbconvert_exporter": "python",
46+
"pygments_lexer": "ipython3",
47+
"version": "3.11.3"
48+
}
49+
},
50+
"nbformat": 4,
51+
"nbformat_minor": 5
52+
}

‎crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"id": "cad32845-44f9-4a53-8b8c-a6b1bb3f3378",
7+
"metadata": {},
8+
"outputs": [],
9+
"source": [
10+
"import math\n",
11+
"\n",
12+
"%matplotlib inline\n",
13+
"\n",
14+
"\n",
15+
"_ = math.pi"
16+
]
17+
},
18+
{
19+
"cell_type": "code",
20+
"execution_count": null,
21+
"id": "d7b8e967-8b4a-493b-b6f7-d5cecfb3a5c3",
22+
"metadata": {},
23+
"outputs": [],
24+
"source": [
25+
"%%timeit\n",
26+
"import sys"
27+
]
28+
}
29+
],
30+
"metadata": {
31+
"kernelspec": {
32+
"display_name": "Python (ruff)",
33+
"language": "python",
34+
"name": "ruff"
35+
},
36+
"language_info": {
37+
"codemirror_mode": {
38+
"name": "ipython",
39+
"version": 3
40+
},
41+
"file_extension": ".py",
42+
"mimetype": "text/x-python",
43+
"name": "python",
44+
"nbconvert_exporter": "python",
45+
"pygments_lexer": "ipython3",
46+
"version": "3.11.3"
47+
}
48+
},
49+
"nbformat": 4,
50+
"nbformat_minor": 5
51+
}

‎crates/ruff/src/autofix/edits.rs

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@
33
use anyhow::{bail, Result};
44

55
use ruff_diagnostics::Edit;
6-
use ruff_python_ast::{self as ast, Arguments, ExceptHandler, Expr, Keyword, Ranged, Stmt};
6+
use ruff_python_ast::{
7+
self as ast, Arguments, ExceptHandler, Expr, Keyword, PySourceType, Ranged, Stmt,
8+
};
79
use ruff_python_codegen::Stylist;
810
use ruff_python_index::Indexer;
9-
use ruff_python_parser::{lexer, Mode};
11+
use ruff_python_parser::{lexer, AsMode};
1012
use ruff_python_trivia::{has_leading_content, is_python_whitespace, PythonWhitespace};
1113
use ruff_source_file::{Locator, NewlineWithTrailingNewline};
1214
use ruff_text_size::{TextLen, TextRange, TextSize};
@@ -88,6 +90,7 @@ pub(crate) fn remove_argument<T: Ranged>(
8890
arguments: &Arguments,
8991
parentheses: Parentheses,
9092
locator: &Locator,
93+
source_type: PySourceType,
9194
) -> Result<Edit> {
9295
// TODO(sbrugman): Preserve trailing comments.
9396
if arguments.keywords.len() + arguments.args.len() > 1 {
@@ -106,7 +109,7 @@ pub(crate) fn remove_argument<T: Ranged>(
106109
let mut seen_comma = false;
107110
for (tok, range) in lexer::lex_starts_at(
108111
locator.slice(arguments.range()),
109-
Mode::Module,
112+
source_type.as_mode(),
110113
arguments.start(),
111114
)
112115
.flatten()
@@ -135,7 +138,7 @@ pub(crate) fn remove_argument<T: Ranged>(
135138
// previous comma to the end of the argument.
136139
for (tok, range) in lexer::lex_starts_at(
137140
locator.slice(arguments.range()),
138-
Mode::Module,
141+
source_type.as_mode(),
139142
arguments.start(),
140143
)
141144
.flatten()

‎crates/ruff/src/checkers/ast/analyze/deferred_scopes.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) {
3737
// Identify any valid runtime imports. If a module is imported at runtime, and
3838
// used at runtime, then by default, we avoid flagging any other
3939
// imports from that model as typing-only.
40-
let enforce_typing_imports = !checker.is_stub
40+
let enforce_typing_imports = !checker.source_type.is_stub()
4141
&& checker.any_enabled(&[
4242
Rule::RuntimeImportInTypeCheckingBlock,
4343
Rule::TypingOnlyFirstPartyImport,
@@ -243,7 +243,7 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) {
243243
pyflakes::rules::unused_annotation(checker, scope, &mut diagnostics);
244244
}
245245

246-
if !checker.is_stub {
246+
if !checker.source_type.is_stub() {
247247
if checker.any_enabled(&[
248248
Rule::UnusedClassMethodArgument,
249249
Rule::UnusedFunctionArgument,

‎crates/ruff/src/checkers/ast/analyze/definitions.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ pub(crate) fn definitions(checker: &mut Checker) {
3030
Rule::MissingTypeKwargs,
3131
Rule::MissingTypeSelf,
3232
]);
33-
let enforce_stubs = checker.is_stub && checker.enabled(Rule::DocstringInStub);
33+
let enforce_stubs = checker.source_type.is_stub() && checker.enabled(Rule::DocstringInStub);
3434
let enforce_stubs_and_runtime = checker.enabled(Rule::IterMethodReturnIterable);
3535
let enforce_docstrings = checker.any_enabled(&[
3636
Rule::BlankLineAfterLastSection,

‎crates/ruff/src/checkers/ast/analyze/expression.rs

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
3131
if let Some(operator) = typing::to_pep604_operator(value, slice, &checker.semantic)
3232
{
3333
if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
34-
if !checker.is_stub
34+
if !checker.source_type.is_stub()
3535
&& checker.settings.target_version < PythonVersion::Py310
3636
&& checker.settings.target_version >= PythonVersion::Py37
3737
&& !checker.semantic.future_annotations()
@@ -44,7 +44,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
4444
}
4545
}
4646
if checker.enabled(Rule::NonPEP604Annotation) {
47-
if checker.is_stub
47+
if checker.source_type.is_stub()
4848
|| checker.settings.target_version >= PythonVersion::Py310
4949
|| (checker.settings.target_version >= PythonVersion::Py37
5050
&& checker.semantic.future_annotations()
@@ -59,7 +59,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
5959

6060
// Ex) list[...]
6161
if checker.enabled(Rule::FutureRequiredTypeAnnotation) {
62-
if !checker.is_stub
62+
if !checker.source_type.is_stub()
6363
&& checker.settings.target_version < PythonVersion::Py39
6464
&& !checker.semantic.future_annotations()
6565
&& checker.semantic.in_annotation()
@@ -176,7 +176,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
176176
typing::to_pep585_generic(expr, &checker.semantic)
177177
{
178178
if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
179-
if !checker.is_stub
179+
if !checker.source_type.is_stub()
180180
&& checker.settings.target_version < PythonVersion::Py39
181181
&& checker.settings.target_version >= PythonVersion::Py37
182182
&& !checker.semantic.future_annotations()
@@ -187,7 +187,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
187187
}
188188
}
189189
if checker.enabled(Rule::NonPEP585Annotation) {
190-
if checker.is_stub
190+
if checker.source_type.is_stub()
191191
|| checker.settings.target_version >= PythonVersion::Py39
192192
|| (checker.settings.target_version >= PythonVersion::Py37
193193
&& checker.semantic.future_annotations()
@@ -272,7 +272,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
272272
]) {
273273
if let Some(replacement) = typing::to_pep585_generic(expr, &checker.semantic) {
274274
if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
275-
if !checker.is_stub
275+
if !checker.source_type.is_stub()
276276
&& checker.settings.target_version < PythonVersion::Py39
277277
&& checker.settings.target_version >= PythonVersion::Py37
278278
&& !checker.semantic.future_annotations()
@@ -285,7 +285,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
285285
}
286286
}
287287
if checker.enabled(Rule::NonPEP585Annotation) {
288-
if checker.is_stub
288+
if checker.source_type.is_stub()
289289
|| checker.settings.target_version >= PythonVersion::Py39
290290
|| (checker.settings.target_version >= PythonVersion::Py37
291291
&& checker.semantic.future_annotations()
@@ -1066,7 +1066,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
10661066
}) => {
10671067
// Ex) `str | None`
10681068
if checker.enabled(Rule::FutureRequiredTypeAnnotation) {
1069-
if !checker.is_stub
1069+
if !checker.source_type.is_stub()
10701070
&& checker.settings.target_version < PythonVersion::Py310
10711071
&& !checker.semantic.future_annotations()
10721072
&& checker.semantic.in_annotation()
@@ -1212,7 +1212,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
12121212
kind: _,
12131213
range: _,
12141214
}) => {
1215-
if checker.is_stub && checker.enabled(Rule::NumericLiteralTooLong) {
1215+
if checker.source_type.is_stub() && checker.enabled(Rule::NumericLiteralTooLong) {
12161216
flake8_pyi::rules::numeric_literal_too_long(checker, expr);
12171217
}
12181218
}
@@ -1221,7 +1221,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
12211221
kind: _,
12221222
range: _,
12231223
}) => {
1224-
if checker.is_stub && checker.enabled(Rule::StringOrBytesTooLong) {
1224+
if checker.source_type.is_stub() && checker.enabled(Rule::StringOrBytesTooLong) {
12251225
flake8_pyi::rules::string_or_bytes_too_long(checker, expr);
12261226
}
12271227
}
@@ -1249,7 +1249,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
12491249
if checker.enabled(Rule::UnicodeKindPrefix) {
12501250
pyupgrade::rules::unicode_kind_prefix(checker, expr, kind.as_deref());
12511251
}
1252-
if checker.is_stub {
1252+
if checker.source_type.is_stub() {
12531253
if checker.enabled(Rule::StringOrBytesTooLong) {
12541254
flake8_pyi::rules::string_or_bytes_too_long(checker, expr);
12551255
}

‎crates/ruff/src/checkers/ast/analyze/parameters.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ pub(crate) fn parameters(parameters: &Parameters, checker: &mut Checker) {
1515
if checker.settings.rules.enabled(Rule::ImplicitOptional) {
1616
ruff::rules::implicit_optional(checker, parameters);
1717
}
18-
if checker.is_stub {
18+
if checker.source_type.is_stub() {
1919
if checker.enabled(Rule::TypedArgumentDefaultInStub) {
2020
flake8_pyi::rules::typed_argument_simple_defaults(checker, parameters);
2121
}

0 commit comments

Comments
 (0)