From 69ee5e4e791a35193d42e8a2a5ff965521aee665 Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Wed, 26 Jul 2023 23:27:22 +0530 Subject: [PATCH] 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 4f98234614fcd6..c6e1913922381a 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 { // 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 97ff354925c4cf..64f820e2c94cfa 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 6d71be0ac85fb2..58da36a4145eed 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 33e9ff317b4512..2bc07b86b4df42 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. @@ -427,7 +433,7 @@ x = 1 let tokens: Vec = 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 a4147e2e7f409f..5dabece32320a8 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 { // 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 1b7c176a4a210a..4b188d6a7b79f6 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 { 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 a9929d12009787..793458b8c94887 100644 --- a/crates/ruff/src/rules/flake8_annotations/rules/definition.rs +++ b/crates/ruff/src/rules/flake8_annotations/rules/definition.rs @@ -706,8 +706,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); @@ -724,8 +729,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 dd78de2573c8e1..308764f8972e28 100644 --- a/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs +++ b/crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs @@ -351,6 +351,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 1ea6663ffdf055..109b3c58b96585 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 { /// ``` /// /// 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 58559ae347349a..cd4182eac38166 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 @@ -66,7 +66,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()) { @@ -78,14 +78,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 { +fn match_parens( + start: TextSize, + locator: &Locator, + is_jupyter_notebook: bool, +) -> Option { 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 d896c07a7eeaa8..1f95ee4db03626 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 04a842f1201f75..a24a2f0d1ae79c 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 4695395b417a81..a787e32e547be7 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 4cad6f78886501..150eb648cd7c14 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>, locator: &Locator, split_on_trailing_comma: bool, + is_jupyter_notebook: bool, ) -> Vec> { 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 32138ca3849196..94fbb2a9922f78 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> { +pub(crate) fn collect_comments<'a>( + range: TextRange, + locator: &'a Locator, + is_jupyter_notebook: bool, +) -> Vec> { 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 54ceb3a8ff5e9c..c1454a10a5830a 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 55412202d84eb8..5587180636e61e 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 fa19caa195500e..543c89decad942 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 { 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 af621cd3e1ad5c..aeeee3647e07ec 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 { // 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 7cb1f8163b5968..e4a4de3c848636 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 + '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 b4d5ca2226ed2e..1b78245083d186 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(location: TextSize, locator: &Locator, f: F) -> Option +fn match_token_before( + location: TextSize, + locator: &Locator, + mode: Mode, + f: F, +) -> Option 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(location: TextSize, locator: &Locator, f: F) -> Option +fn match_token_after( + location: TextSize, + locator: &Locator, + mode: Mode, + f: F, +) -> Option 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(location: TextSize, locator: &Locator, f: F) -> Option +fn match_token_or_closing_brace( + location: TextSize, + locator: &Locator, + mode: Mode, + f: F, +) -> Option 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 { + 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 d1bd29c2e97947..ff7eb99a4397a1 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 = 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 a82fc6760a245e..04492d26bd0679 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 = 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 7ee67310dfea11..feaeafb51f8152 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 { +fn create_remove_param_fix( + locator: &Locator, + expr: &Expr, + mode_param: &Expr, + is_jupyter_notebook: bool, +) -> Result { 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 = 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 d32f8fb7ab8cac..d89331d462d8f4 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 { let (first, second) = if stdout.start() < stderr.start() { (stdout, stderr) @@ -74,6 +75,7 @@ fn generate_fix( args, keywords, false, + is_jupyter_notebook, )?], )) } @@ -121,7 +123,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 3e85b88d8c3fba..ea1ab1a56336d3 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 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 ea26a66e0f92e3..1264eac3135b55 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 5ca5c0be574173..f1e3f7516a7521 100644 --- a/crates/ruff_python_parser/src/lib.rs +++ b/crates/ruff_python_parser/src/lib.rs @@ -39,9 +39,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 { +pub fn first_colon_range( + range: TextRange, + source: &str, + is_jupyter_notebook: bool, +) -> Option { 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); @@ -163,6 +172,7 @@ mod tests { let range = first_colon_range( TextRange::new(TextSize::from(0), contents.text_len()), contents, + false, ) .unwrap(); assert_eq!(&contents[range], ":");