Implement isort custom sections and ordering (astral-sh#2419)

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

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

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,27 +70,42 @@ 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
             .get(&target_version.as_tuple())
             .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<ImportType, ImportBlock<'a>> {
-    let mut block_by_type: BTreeMap<ImportType, ImportBlock> = BTreeMap::default();
+) -> BTreeMap<ImportSection, ImportBlock<'a>> {
+    let mut block_by_type: BTreeMap<ImportSection, ImportBlock> = 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<String>,
     /// A set of user-provided standard library modules.
     pub standard_library: BTreeSet<String>,
+    /// A map of additional user-provided sections.
+    pub sections: FxHashMap<String, BTreeSet<String>>,
     /// 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<String>,
         local_folder: Vec<String>,
         standard_library: Vec<String>,
+        sections: FxHashMap<String, Vec<String>>,
     ) -> 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().flatten())
             .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
         }
     }
 }