diff --git a/crates/ruff/resources/test/fixtures/isort/sections.py b/crates/ruff/resources/test/fixtures/isort/sections.py new file mode 100644 index 00000000000000..dffdeddce6b628 --- /dev/null +++ b/crates/ruff/resources/test/fixtures/isort/sections.py @@ -0,0 +1,7 @@ +from __future__ import annotations +import os +import sys +import pytz +import django.settings +from library import foo +from . import local 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 7c13715c0c7070..8af8156563a9f0 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 @@ -6,7 +6,7 @@ use ruff_python_semantic::binding::{ Binding, BindingKind, ExecutionContext, FromImportation, Importation, SubmoduleImportation, }; -use crate::rules::isort::{categorize, ImportType}; +use crate::rules::isort::{categorize, ImportSection, ImportType}; use crate::settings::Settings; /// ## What it does @@ -294,25 +294,31 @@ pub fn typing_only_runtime_import( &settings.isort.known_modules, settings.target_version, ) { - ImportType::LocalFolder | ImportType::FirstParty => Some(Diagnostic::new( - TypingOnlyFirstPartyImport { - full_name: full_name.to_string(), - }, - binding.range, - )), - ImportType::ThirdParty => Some(Diagnostic::new( - TypingOnlyThirdPartyImport { - full_name: full_name.to_string(), - }, - binding.range, - )), - ImportType::StandardLibrary => Some(Diagnostic::new( + ImportSection::Known(ImportType::LocalFolder | ImportType::FirstParty) => { + Some(Diagnostic::new( + TypingOnlyFirstPartyImport { + full_name: full_name.to_string(), + }, + binding.range, + )) + } + ImportSection::Known(ImportType::ThirdParty) | ImportSection::UserDefined(_) => { + Some(Diagnostic::new( + TypingOnlyThirdPartyImport { + full_name: full_name.to_string(), + }, + binding.range, + )) + } + ImportSection::Known(ImportType::StandardLibrary) => Some(Diagnostic::new( TypingOnlyStandardLibraryImport { full_name: full_name.to_string(), }, binding.range, )), - ImportType::Future => unreachable!("`__future__` imports should be marked as used"), + ImportSection::Known(ImportType::Future) => { + unreachable!("`__future__` imports should be marked as used") + } } } else { None diff --git a/crates/ruff/src/rules/isort/categorize.rs b/crates/ruff/src/rules/isort/categorize.rs index 424431fa008bcc..6303af310c53a7 100644 --- a/crates/ruff/src/rules/isort/categorize.rs +++ b/crates/ruff/src/rules/isort/categorize.rs @@ -3,6 +3,7 @@ use std::path::{Path, PathBuf}; use std::{fs, iter}; use log::debug; +use rustc_hash::FxHashMap; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use strum_macros::EnumIter; @@ -37,6 +38,15 @@ pub enum ImportType { LocalFolder, } +#[derive( + Debug, PartialOrd, Ord, PartialEq, Eq, Clone, Serialize, Deserialize, JsonSchema, CacheKey, +)] +#[serde(untagged)] +pub enum ImportSection { + Known(ImportType), + UserDefined(String), +} + #[derive(Debug)] enum Reason<'a> { NonZeroLevel, @@ -49,6 +59,7 @@ enum Reason<'a> { SamePackage, SourceMatch(&'a Path), NoMatch, + UserDefinedSection, } #[allow(clippy::too_many_arguments)] @@ -59,13 +70,16 @@ pub fn categorize( package: Option<&Path>, known_modules: &KnownModules, target_version: PythonVersion, -) -> ImportType { +) -> ImportSection { let module_base = module_name.split('.').next().unwrap(); let (import_type, reason) = { if level.map_or(false, |level| level > 0) { - (ImportType::LocalFolder, Reason::NonZeroLevel) + ( + ImportSection::Known(ImportType::LocalFolder), + Reason::NonZeroLevel, + ) } else if module_base == "__future__" { - (ImportType::Future, Reason::Future) + (ImportSection::Known(ImportType::Future), Reason::Future) } else if let Some((import_type, reason)) = known_modules.categorize(module_name) { (import_type, reason) } else if KNOWN_STANDARD_LIBRARY @@ -73,13 +87,25 @@ pub fn categorize( .unwrap() .contains(module_base) { - (ImportType::StandardLibrary, Reason::KnownStandardLibrary) + ( + ImportSection::Known(ImportType::StandardLibrary), + Reason::KnownStandardLibrary, + ) } else if same_package(package, module_base) { - (ImportType::FirstParty, Reason::SamePackage) + ( + ImportSection::Known(ImportType::FirstParty), + Reason::SamePackage, + ) } else if let Some(src) = match_sources(src, module_base) { - (ImportType::FirstParty, Reason::SourceMatch(src)) + ( + ImportSection::Known(ImportType::FirstParty), + Reason::SourceMatch(src), + ) } else { - (ImportType::ThirdParty, Reason::NoMatch) + ( + ImportSection::Known(ImportType::ThirdParty), + Reason::NoMatch, + ) } }; debug!( @@ -116,8 +142,8 @@ pub fn categorize_imports<'a>( package: Option<&Path>, known_modules: &KnownModules, target_version: PythonVersion, -) -> BTreeMap> { - let mut block_by_type: BTreeMap = BTreeMap::default(); +) -> BTreeMap> { + let mut block_by_type: BTreeMap = BTreeMap::default(); // Categorize `StmtKind::Import`. for (alias, comments) in block.import { let import_type = categorize( @@ -195,6 +221,8 @@ pub struct KnownModules { pub local_folder: BTreeSet, /// A set of user-provided standard library modules. pub standard_library: BTreeSet, + /// A map of additional user-provided sections. + pub sections: FxHashMap>, /// Whether any of the known modules are submodules (e.g., `foo.bar`, as opposed to `foo`). has_submodules: bool, } @@ -205,29 +233,36 @@ impl KnownModules { third_party: Vec, local_folder: Vec, standard_library: Vec, + sections: FxHashMap>, ) -> Self { let first_party = BTreeSet::from_iter(first_party); let third_party = BTreeSet::from_iter(third_party); let local_folder = BTreeSet::from_iter(local_folder); let standard_library = BTreeSet::from_iter(standard_library); + let sections: FxHashMap<_, _> = sections + .into_iter() + .map(|(k, v)| (k, BTreeSet::from_iter(v))) + .collect(); let has_submodules = first_party .iter() .chain(third_party.iter()) .chain(local_folder.iter()) .chain(standard_library.iter()) + .chain(sections.values().flat_map(std::collections::BTreeSet::iter)) .any(|m| m.contains('.')); Self { first_party, third_party, local_folder, standard_library, + sections, has_submodules, } } - /// Return the [`ImportType`] for a given module, if it's been categorized as a known module + /// Return the [`ImportSection`] for a given module, if it's been categorized as a known module /// by the user. - fn categorize(&self, module_name: &str) -> Option<(ImportType, Reason)> { + fn categorize(&self, module_name: &str) -> Option<(ImportSection, Reason)> { if self.has_submodules { // Check all module prefixes from the longest to the shortest (e.g., given // `foo.bar.baz`, check `foo.bar.baz`, then `foo.bar`, then `foo`, taking the first, @@ -239,34 +274,46 @@ impl KnownModules { .rev() { let submodule = &module_name[0..i]; - if self.first_party.contains(submodule) { - return Some((ImportType::FirstParty, Reason::KnownFirstParty)); - } - if self.third_party.contains(submodule) { - return Some((ImportType::ThirdParty, Reason::KnownThirdParty)); - } - if self.local_folder.contains(submodule) { - return Some((ImportType::LocalFolder, Reason::KnownLocalFolder)); - } - if self.standard_library.contains(submodule) { - return Some((ImportType::StandardLibrary, Reason::ExtraStandardLibrary)); + if let Some(result) = self.categorize_submodule(submodule) { + return Some(result); } } None } else { // Happy path: no submodules, so we can check the module base and be done. let module_base = module_name.split('.').next().unwrap(); - if self.first_party.contains(module_base) { - Some((ImportType::FirstParty, Reason::KnownFirstParty)) - } else if self.third_party.contains(module_base) { - Some((ImportType::ThirdParty, Reason::KnownThirdParty)) - } else if self.local_folder.contains(module_base) { - Some((ImportType::LocalFolder, Reason::KnownLocalFolder)) - } else if self.standard_library.contains(module_base) { - Some((ImportType::StandardLibrary, Reason::ExtraStandardLibrary)) - } else { - None - } + self.categorize_submodule(module_base) + } + } + + fn categorize_submodule(&self, submodule: &str) -> Option<(ImportSection, Reason)> { + if let Some((section, _)) = self.sections.iter().find(|(_, v)| v.contains(submodule)) { + Some(( + ImportSection::UserDefined(section.clone()), + Reason::UserDefinedSection, + )) + } else if self.first_party.contains(submodule) { + Some(( + ImportSection::Known(ImportType::FirstParty), + Reason::KnownFirstParty, + )) + } else if self.third_party.contains(submodule) { + Some(( + ImportSection::Known(ImportType::ThirdParty), + Reason::KnownThirdParty, + )) + } else if self.local_folder.contains(submodule) { + Some(( + ImportSection::Known(ImportType::LocalFolder), + Reason::KnownLocalFolder, + )) + } else if self.standard_library.contains(submodule) { + Some(( + ImportSection::Known(ImportType::StandardLibrary), + Reason::ExtraStandardLibrary, + )) + } else { + None } } } diff --git a/crates/ruff/src/rules/isort/mod.rs b/crates/ruff/src/rules/isort/mod.rs index a862abfd60ebd3..6207b96943ebf5 100644 --- a/crates/ruff/src/rules/isort/mod.rs +++ b/crates/ruff/src/rules/isort/mod.rs @@ -9,7 +9,7 @@ use strum::IntoEnumIterator; use crate::rules::isort::categorize::KnownModules; use annotate::annotate_imports; use categorize::categorize_imports; -pub use categorize::{categorize, ImportType}; +pub use categorize::{categorize, ImportSection, ImportType}; use comments::Comment; use normalize::normalize_imports; use order::order_imports; @@ -131,11 +131,12 @@ pub fn format_imports( classes: &BTreeSet, constants: &BTreeSet, variables: &BTreeSet, - no_lines_before: &BTreeSet, + no_lines_before: &BTreeSet, lines_after_imports: isize, lines_between_types: usize, forced_separate: &[String], target_version: PythonVersion, + section_order: &[ImportSection], ) -> String { let trailer = &block.trailer; let block = annotate_imports(&block.imports, comments, locator, split_on_trailing_comma); @@ -143,6 +144,23 @@ pub fn format_imports( // Normalize imports (i.e., deduplicate, aggregate `from` imports). let block = normalize_imports(block, combine_as_imports, force_single_line); + // Make sure all sections (built-in and user-defined) are present in the section order. + let mut section_order: Vec<_> = section_order.to_vec(); + for known_section in ImportType::iter().map(ImportSection::Known) { + if !section_order.contains(&known_section) { + section_order.push(known_section); + } + } + for section in known_modules + .sections + .keys() + .map(|section| ImportSection::UserDefined(section.clone())) + { + if !section_order.contains(§ion) { + section_order.push(section); + } + } + let mut output = String::new(); for block in split::split_by_forced_separate(block, forced_separate) { @@ -167,6 +185,7 @@ pub fn format_imports( no_lines_before, lines_between_types, target_version, + §ion_order, ); if !block_output.is_empty() && !output.is_empty() { @@ -221,9 +240,10 @@ fn format_import_block( classes: &BTreeSet, constants: &BTreeSet, variables: &BTreeSet, - no_lines_before: &BTreeSet, + no_lines_before: &BTreeSet, lines_between_types: usize, target_version: PythonVersion, + section_order: &[ImportSection], ) -> String { // Categorize by type (e.g., first-party vs. third-party). let mut block_by_type = categorize_imports(block, src, package, known_modules, target_version); @@ -233,10 +253,10 @@ fn format_import_block( // Generate replacement source code. let mut is_first_block = true; let mut pending_lines_before = false; - for import_type in ImportType::iter() { - let import_block = block_by_type.remove(&import_type); + for import_section in section_order { + let import_block = block_by_type.remove(import_section); - if !no_lines_before.contains(&import_type) { + if !no_lines_before.contains(import_section) { pending_lines_before = true; } let Some(import_block) = import_block else { @@ -332,10 +352,11 @@ mod tests { use anyhow::Result; use insta::assert_yaml_snapshot; + use rustc_hash::FxHashMap; use test_case::test_case; use crate::registry::Rule; - use crate::rules::isort::categorize::KnownModules; + use crate::rules::isort::categorize::{ImportSection, KnownModules}; use crate::settings::Settings; use crate::test::{test_path, test_resource_path}; @@ -412,6 +433,7 @@ mod tests { vec!["foo".to_string(), "__future__".to_string()], vec![], vec![], + FxHashMap::default(), ), ..super::settings::Settings::default() }, @@ -435,6 +457,7 @@ mod tests { vec!["foo.bar".to_string()], vec![], vec![], + FxHashMap::default(), ), ..super::settings::Settings::default() }, @@ -476,6 +499,7 @@ mod tests { vec![], vec!["ruff".to_string()], vec![], + FxHashMap::default(), ), ..super::settings::Settings::default() }, @@ -826,11 +850,11 @@ mod tests { &Settings { isort: super::settings::Settings { no_lines_before: BTreeSet::from([ - ImportType::Future, - ImportType::StandardLibrary, - ImportType::ThirdParty, - ImportType::FirstParty, - ImportType::LocalFolder, + ImportSection::Known(ImportType::Future), + ImportSection::Known(ImportType::StandardLibrary), + ImportSection::Known(ImportType::ThirdParty), + ImportSection::Known(ImportType::FirstParty), + ImportSection::Known(ImportType::LocalFolder), ]), ..super::settings::Settings::default() }, @@ -854,8 +878,8 @@ mod tests { &Settings { isort: super::settings::Settings { no_lines_before: BTreeSet::from([ - ImportType::StandardLibrary, - ImportType::LocalFolder, + ImportSection::Known(ImportType::StandardLibrary), + ImportSection::Known(ImportType::LocalFolder), ]), ..super::settings::Settings::default() }, @@ -930,4 +954,60 @@ mod tests { assert_yaml_snapshot!(snapshot, diagnostics); Ok(()) } + + #[test_case(Path::new("sections.py"))] + fn sections(path: &Path) -> Result<()> { + let snapshot = format!("sections_{}", path.to_string_lossy()); + let diagnostics = test_path( + Path::new("isort").join(path).as_path(), + &Settings { + src: vec![test_resource_path("fixtures/isort")], + isort: super::settings::Settings { + known_modules: KnownModules::new( + vec![], + vec![], + vec![], + vec![], + FxHashMap::from_iter([("django".to_string(), vec!["django".to_string()])]), + ), + ..super::settings::Settings::default() + }, + ..Settings::for_rule(Rule::UnsortedImports) + }, + )?; + assert_yaml_snapshot!(snapshot, diagnostics); + Ok(()) + } + + #[test_case(Path::new("sections.py"))] + fn section_order(path: &Path) -> Result<()> { + let snapshot = format!("section_order_{}", path.to_string_lossy()); + let diagnostics = test_path( + Path::new("isort").join(path).as_path(), + &Settings { + src: vec![test_resource_path("fixtures/isort")], + isort: super::settings::Settings { + known_modules: KnownModules::new( + vec!["library".to_string()], + vec![], + vec![], + vec![], + FxHashMap::from_iter([("django".to_string(), vec!["django".to_string()])]), + ), + section_order: vec![ + ImportSection::Known(ImportType::Future), + ImportSection::Known(ImportType::StandardLibrary), + ImportSection::Known(ImportType::ThirdParty), + ImportSection::UserDefined("django".to_string()), + ImportSection::Known(ImportType::FirstParty), + ImportSection::Known(ImportType::LocalFolder), + ], + ..super::settings::Settings::default() + }, + ..Settings::for_rule(Rule::UnsortedImports) + }, + )?; + assert_yaml_snapshot!(snapshot, diagnostics); + Ok(()) + } } diff --git a/crates/ruff/src/rules/isort/rules/organize_imports.rs b/crates/ruff/src/rules/isort/rules/organize_imports.rs index ddbde6ed69d881..b5e1e53edefe52 100644 --- a/crates/ruff/src/rules/isort/rules/organize_imports.rs +++ b/crates/ruff/src/rules/isort/rules/organize_imports.rs @@ -139,6 +139,7 @@ pub fn organize_imports( settings.isort.lines_between_types, &settings.isort.forced_separate, settings.target_version, + &settings.isort.section_order, ); // Expand the span the entire range, including leading and trailing space. diff --git a/crates/ruff/src/rules/isort/settings.rs b/crates/ruff/src/rules/isort/settings.rs index c9acc603a6e1c6..f7ef72591b4b96 100644 --- a/crates/ruff/src/rules/isort/settings.rs +++ b/crates/ruff/src/rules/isort/settings.rs @@ -1,5 +1,6 @@ //! Settings for the `isort` plugin. +use rustc_hash::FxHashMap; use std::collections::BTreeSet; use schemars::JsonSchema; @@ -8,7 +9,7 @@ use serde::{Deserialize, Serialize}; use crate::rules::isort::categorize::KnownModules; use ruff_macros::{CacheKey, ConfigurationOptions}; -use super::categorize::ImportType; +use super::categorize::ImportSection; #[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize, CacheKey, JsonSchema)] #[serde(deny_unknown_fields, rename_all = "kebab-case")] @@ -226,14 +227,14 @@ pub struct Options { pub variables: Option>, #[option( default = r#"[]"#, - value_type = r#"list["future" | "standard-library" | "third-party" | "first-party" | "local-folder"]"#, + value_type = r#"list[str]"#, example = r#" no-lines-before = ["future", "standard-library"] "# )] /// A list of sections that should _not_ be delineated from the previous /// section via empty lines. - pub no_lines_before: Option>, + pub no_lines_before: Option>, #[option( default = r#"-1"#, value_type = "int", @@ -265,6 +266,28 @@ pub struct Options { /// A list of modules to separate into auxiliary block(s) of imports, /// in the order specified. pub forced_separate: Option>, + #[option( + default = r#"[]"#, + value_type = "list[str]", + example = r#" + section-order = ["future", "standard-library", "first-party", "local-folder", "third-party"] + "# + )] + /// Override in which order the sections should be output. Can be used to move custom sections. + pub section_order: Option>, + // Tables are required to go last. + #[option( + default = "{}", + value_type = "dict[str, list[str]]", + example = r#" + # Group all Django imports into a separate section. + [tool.ruff.isort.sections] + "django" = ["django"] + "# + )] + /// A list of mappings from section names to modules. + /// By default custom sections are output last, but this can be overridden with `section-order`. + pub sections: Option>>, } #[derive(Debug, CacheKey)] @@ -284,10 +307,11 @@ pub struct Settings { pub classes: BTreeSet, pub constants: BTreeSet, pub variables: BTreeSet, - pub no_lines_before: BTreeSet, + pub no_lines_before: BTreeSet, pub lines_after_imports: isize, pub lines_between_types: usize, pub forced_separate: Vec, + pub section_order: Vec, } impl Default for Settings { @@ -311,6 +335,7 @@ impl Default for Settings { lines_after_imports: -1, lines_between_types: 0, forced_separate: Vec::new(), + section_order: Vec::new(), } } } @@ -329,6 +354,7 @@ impl From for Settings { options.known_third_party.unwrap_or_default(), options.known_local_folder.unwrap_or_default(), options.extra_standard_library.unwrap_or_default(), + options.sections.unwrap_or_default(), ), order_by_type: options.order_by_type.unwrap_or(true), relative_imports_order: options.relative_imports_order.unwrap_or_default(), @@ -343,6 +369,7 @@ impl From for Settings { lines_after_imports: options.lines_after_imports.unwrap_or(-1), lines_between_types: options.lines_between_types.unwrap_or_default(), forced_separate: Vec::from_iter(options.forced_separate.unwrap_or_default()), + section_order: Vec::from_iter(options.section_order.unwrap_or_default()), } } } @@ -377,6 +404,15 @@ impl From for Options { lines_after_imports: Some(settings.lines_after_imports), lines_between_types: Some(settings.lines_between_types), forced_separate: Some(settings.forced_separate.into_iter().collect()), + section_order: Some(settings.section_order.into_iter().collect()), + sections: Some( + settings + .known_modules + .sections + .into_iter() + .map(|(k, v)| (k, v.into_iter().collect())) + .collect(), + ), } } } diff --git a/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__section_order_sections.py.snap b/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__section_order_sections.py.snap new file mode 100644 index 00000000000000..8f47cc75b208ad --- /dev/null +++ b/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__section_order_sections.py.snap @@ -0,0 +1,26 @@ +--- +source: crates/ruff/src/rules/isort/mod.rs +expression: diagnostics +--- +- kind: + name: UnsortedImports + body: Import block is un-sorted or un-formatted + suggestion: Organize imports + fixable: true + location: + row: 1 + column: 0 + end_location: + row: 8 + column: 0 + fix: + edits: + - content: "from __future__ import annotations\n\nimport os\nimport sys\n\nimport pytz\n\nimport django.settings\n\nfrom library import foo\n\nfrom . import local\n" + location: + row: 1 + column: 0 + end_location: + row: 8 + column: 0 + parent: ~ + diff --git a/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__sections_sections.py.snap b/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__sections_sections.py.snap new file mode 100644 index 00000000000000..4793d77472d91a --- /dev/null +++ b/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__sections_sections.py.snap @@ -0,0 +1,26 @@ +--- +source: crates/ruff/src/rules/isort/mod.rs +expression: diagnostics +--- +- kind: + name: UnsortedImports + body: Import block is un-sorted or un-formatted + suggestion: Organize imports + fixable: true + location: + row: 1 + column: 0 + end_location: + row: 8 + column: 0 + fix: + edits: + - content: "from __future__ import annotations\n\nimport os\nimport sys\n\nimport pytz\nfrom library import foo\n\nfrom . import local\n\nimport django.settings\n" + location: + row: 1 + column: 0 + end_location: + row: 8 + column: 0 + parent: ~ + diff --git a/ruff.schema.json b/ruff.schema.json index aaa04c3f6c008c..ee3051dc1082c4 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -986,6 +986,16 @@ }, "additionalProperties": false }, + "ImportSection": { + "anyOf": [ + { + "$ref": "#/definitions/ImportType" + }, + { + "type": "string" + } + ] + }, "ImportType": { "type": "string", "enum": [ @@ -1131,7 +1141,7 @@ "null" ], "items": { - "$ref": "#/definitions/ImportType" + "$ref": "#/definitions/ImportSection" } }, "order-by-type": { @@ -1162,6 +1172,29 @@ "type": "string" } }, + "section-order": { + "description": "Override in which order the sections should be output. Can be used to move custom sections.", + "type": [ + "array", + "null" + ], + "items": { + "$ref": "#/definitions/ImportSection" + } + }, + "sections": { + "description": "A list of mappings from section names to modules. By default custom sections are output last, but this can be overridden with `section-order`.", + "type": [ + "object", + "null" + ], + "additionalProperties": { + "type": "array", + "items": { + "type": "string" + } + } + }, "single-line-exclusions": { "description": "One or more modules to exclude from the single line rule.", "type": [