[`isort`]: support submodules in known_(first|third)_party config options (#3768)

crates/ruff/resources/test/fixtures/isort/separate_subpackage_first_and_third_party_imports.py

@@ -0,0 +1,8 @@
+import sys
+import baz
+from foo import bar, baz
+from foo.bar import blah, blub
+from foo.bar.baz import something
+import foo
+import foo.bar
+import foo.bar.baz

crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs

@@ -178,19 +178,15 @@ pub fn typing_only_runtime_import(
         // Extract the module base and level from the full name.
         // Ex) `foo.bar.baz` -> `foo`, `0`
         // Ex) `.foo.bar.baz` -> `foo`, `1`
-        let module_base = full_name.split('.').next().unwrap();
         let level = full_name.chars().take_while(|c| *c == '.').count();
 
         // Categorize the import.
         match categorize(
-            module_base,
+            full_name,
             Some(&level),
             &settings.src,
             package,
-            &settings.isort.known_first_party,
+            &settings.isort.known_modules,
-            &settings.isort.known_third_party,
-            &settings.isort.known_local_folder,
-            &settings.isort.extra_standard_library,
             settings.target_version,
         ) {
             ImportType::LocalFolder | ImportType::FirstParty => Some(Diagnostic::new(

crates/ruff/src/rules/isort/categorize.rs

@@ -1,6 +1,6 @@
 use std::collections::{BTreeMap, BTreeSet};
-use std::fs;
 use std::path::{Path, PathBuf};
+use std::{fs, iter};
 
 use log::debug;
 use schemars::JsonSchema;
@@ -52,29 +52,21 @@ enum Reason<'a> {
 
 #[allow(clippy::too_many_arguments)]
 pub fn categorize(
-    module_base: &str,
+    module_name: &str,
     level: Option<&usize>,
     src: &[PathBuf],
     package: Option<&Path>,
-    known_first_party: &BTreeSet<String>,
+    known_modules: &KnownModules,
-    known_third_party: &BTreeSet<String>,
-    known_local_folder: &BTreeSet<String>,
-    extra_standard_library: &BTreeSet<String>,
     target_version: PythonVersion,
 ) -> ImportType {
+    let module_base = module_name.split('.').next().unwrap();
     let (import_type, reason) = {
         if level.map_or(false, |level| *level > 0) {
             (ImportType::LocalFolder, Reason::NonZeroLevel)
-        } else if known_first_party.contains(module_base) {
-            (ImportType::FirstParty, Reason::KnownFirstParty)
-        } else if known_third_party.contains(module_base) {
-            (ImportType::ThirdParty, Reason::KnownThirdParty)
-        } else if known_local_folder.contains(module_base) {
-            (ImportType::LocalFolder, Reason::KnownLocalFolder)
-        } else if extra_standard_library.contains(module_base) {
-            (ImportType::StandardLibrary, Reason::ExtraStandardLibrary)
         } else if module_base == "__future__" {
             (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()
@@ -91,7 +83,7 @@ pub fn categorize(
     };
     debug!(
         "Categorized '{}' as {:?} ({:?})",
-        module_base, import_type, reason
+        module_name, import_type, reason
     );
     import_type
 }
@@ -121,24 +113,18 @@ pub fn categorize_imports<'a>(
     block: ImportBlock<'a>,
     src: &[PathBuf],
     package: Option<&Path>,
-    known_first_party: &BTreeSet<String>,
+    known_modules: &KnownModules,
-    known_third_party: &BTreeSet<String>,
-    known_local_folder: &BTreeSet<String>,
-    extra_standard_library: &BTreeSet<String>,
     target_version: PythonVersion,
 ) -> BTreeMap<ImportType, ImportBlock<'a>> {
     let mut block_by_type: BTreeMap<ImportType, ImportBlock> = BTreeMap::default();
     // Categorize `StmtKind::Import`.
     for (alias, comments) in block.import {
         let import_type = categorize(
-            &alias.module_base(),
+            &alias.module_name(),
             None,
             src,
             package,
-            known_first_party,
+            known_modules,
-            known_third_party,
-            known_local_folder,
-            extra_standard_library,
             target_version,
         );
         block_by_type
@@ -150,14 +136,11 @@ pub fn categorize_imports<'a>(
     // Categorize `StmtKind::ImportFrom` (without re-export).
     for (import_from, aliases) in block.import_from {
         let classification = categorize(
-            &import_from.module_base(),
+            &import_from.module_name(),
             import_from.level,
             src,
             package,
-            known_first_party,
+            known_modules,
-            known_third_party,
-            known_local_folder,
-            extra_standard_library,
             target_version,
         );
         block_by_type
@@ -169,14 +152,11 @@ pub fn categorize_imports<'a>(
     // Categorize `StmtKind::ImportFrom` (with re-export).
     for ((import_from, alias), aliases) in block.import_from_as {
         let classification = categorize(
-            &import_from.module_base(),
+            &import_from.module_name(),
             import_from.level,
             src,
             package,
-            known_first_party,
+            known_modules,
-            known_third_party,
-            known_local_folder,
-            extra_standard_library,
             target_version,
         );
         block_by_type
@@ -188,14 +168,11 @@ pub fn categorize_imports<'a>(
     // Categorize `StmtKind::ImportFrom` (with star).
     for (import_from, comments) in block.import_from_star {
         let classification = categorize(
-            &import_from.module_base(),
+            &import_from.module_name(),
             import_from.level,
             src,
             package,
-            known_first_party,
+            known_modules,
-            known_third_party,
-            known_local_folder,
-            extra_standard_library,
             target_version,
         );
         block_by_type
@@ -206,3 +183,89 @@ pub fn categorize_imports<'a>(
     }
     block_by_type
 }
+
+#[derive(Debug, Default, CacheKey)]
+pub struct KnownModules {
+    /// A set of user-provided first-party modules.
+    pub first_party: BTreeSet<String>,
+    /// A set of user-provided third-party modules.
+    pub third_party: BTreeSet<String>,
+    /// A set of user-provided local folder modules.
+    pub local_folder: BTreeSet<String>,
+    /// A set of user-provided standard library modules.
+    pub standard_library: BTreeSet<String>,
+    /// Whether any of the known modules are submodules (e.g., `foo.bar`, as opposed to `foo`).
+    has_submodules: bool,
+}
+
+impl KnownModules {
+    pub fn new(
+        first_party: Vec<String>,
+        third_party: Vec<String>,
+        local_folder: Vec<String>,
+        standard_library: 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 has_submodules = first_party
+            .iter()
+            .chain(third_party.iter())
+            .chain(local_folder.iter())
+            .chain(standard_library.iter())
+            .any(|m| m.contains('.'));
+        Self {
+            first_party,
+            third_party,
+            local_folder,
+            standard_library,
+            has_submodules,
+        }
+    }
+
+    /// Return the [`ImportType`] 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)> {
+        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,
+            // most precise match).
+            for i in module_name
+                .match_indices('.')
+                .map(|(i, _)| i)
+                .chain(iter::once(module_name.len()))
+                .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));
+                }
+            }
+            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
+            }
+        }
+    }
+}

crates/ruff/src/rules/isort/mod.rs

@@ -6,6 +6,7 @@ use std::path::{Path, PathBuf};
 use itertools::Either::{Left, Right};
 use strum::IntoEnumIterator;
 
+use crate::rules::isort::categorize::KnownModules;
 use annotate::annotate_imports;
 use categorize::categorize_imports;
 pub use categorize::{categorize, ImportType};
@@ -118,14 +119,11 @@ pub fn format_imports(
     src: &[PathBuf],
     package: Option<&Path>,
     combine_as_imports: bool,
-    extra_standard_library: &BTreeSet<String>,
     force_single_line: bool,
     force_sort_within_sections: bool,
     force_wrap_aliases: bool,
     force_to_top: &BTreeSet<String>,
-    known_first_party: &BTreeSet<String>,
+    known_modules: &KnownModules,
-    known_third_party: &BTreeSet<String>,
-    known_local_folder: &BTreeSet<String>,
     order_by_type: bool,
     relative_imports_order: RelativeImportsOrder,
     single_line_exclusions: &BTreeSet<String>,
@@ -154,14 +152,11 @@ pub fn format_imports(
             stylist,
             src,
             package,
-            extra_standard_library,
             force_single_line,
             force_sort_within_sections,
             force_wrap_aliases,
             force_to_top,
-            known_first_party,
+            known_modules,
-            known_third_party,
-            known_local_folder,
             order_by_type,
             relative_imports_order,
             single_line_exclusions,
@@ -214,14 +209,11 @@ fn format_import_block(
     stylist: &Stylist,
     src: &[PathBuf],
     package: Option<&Path>,
-    extra_standard_library: &BTreeSet<String>,
     force_single_line: bool,
     force_sort_within_sections: bool,
     force_wrap_aliases: bool,
     force_to_top: &BTreeSet<String>,
-    known_first_party: &BTreeSet<String>,
+    known_modules: &KnownModules,
-    known_third_party: &BTreeSet<String>,
-    known_local_folder: &BTreeSet<String>,
     order_by_type: bool,
     relative_imports_order: RelativeImportsOrder,
     single_line_exclusions: &BTreeSet<String>,
@@ -234,16 +226,7 @@ fn format_import_block(
     target_version: PythonVersion,
 ) -> String {
     // Categorize by type (e.g., first-party vs. third-party).
-    let mut block_by_type = categorize_imports(
+    let mut block_by_type = categorize_imports(block, src, package, known_modules, target_version);
-        block,
-        src,
-        package,
-        known_first_party,
-        known_third_party,
-        known_local_folder,
-        extra_standard_library,
-        target_version,
-    );
 
     let mut output = String::new();
 
@@ -352,6 +335,7 @@ mod tests {
     use test_case::test_case;
 
     use crate::registry::Rule;
+    use crate::rules::isort::categorize::KnownModules;
     use crate::settings::Settings;
     use crate::test::{test_path, test_resource_path};
 
@@ -416,6 +400,52 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(Path::new("separate_subpackage_first_and_third_party_imports.py"))]
+    fn separate_modules(path: &Path) -> Result<()> {
+        let snapshot = format!("1_{}", path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("isort").join(path).as_path(),
+            &Settings {
+                isort: super::settings::Settings {
+                    known_modules: KnownModules::new(
+                        vec!["foo.bar".to_string(), "baz".to_string()],
+                        vec!["foo".to_string(), "__future__".to_string()],
+                        vec![],
+                        vec![],
+                    ),
+                    ..super::settings::Settings::default()
+                },
+                src: vec![test_resource_path("fixtures/isort")],
+                ..Settings::for_rule(Rule::UnsortedImports)
+            },
+        )?;
+        assert_yaml_snapshot!(snapshot, diagnostics);
+        Ok(())
+    }
+
+    #[test_case(Path::new("separate_subpackage_first_and_third_party_imports.py"))]
+    fn separate_modules_first_party(path: &Path) -> Result<()> {
+        let snapshot = format!("2_{}", path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("isort").join(path).as_path(),
+            &Settings {
+                isort: super::settings::Settings {
+                    known_modules: KnownModules::new(
+                        vec!["foo".to_string()],
+                        vec!["foo.bar".to_string()],
+                        vec![],
+                        vec![],
+                    ),
+                    ..super::settings::Settings::default()
+                },
+                src: vec![test_resource_path("fixtures/isort")],
+                ..Settings::for_rule(Rule::UnsortedImports)
+            },
+        )?;
+        assert_yaml_snapshot!(snapshot, diagnostics);
+        Ok(())
+    }
+
     // Test currently disabled as line endings are automatically converted to
     // platform-appropriate ones in CI/CD #[test_case(Path::new("
     // line_ending_crlf.py"))] #[test_case(Path::new("line_ending_lf.py"))]
@@ -441,7 +471,12 @@ mod tests {
             Path::new("isort").join(path).as_path(),
             &Settings {
                 isort: super::settings::Settings {
-                    known_local_folder: BTreeSet::from(["ruff".to_string()]),
+                    known_modules: KnownModules::new(
+                        vec![],
+                        vec![],
+                        vec!["ruff".to_string()],
+                        vec![],
+                    ),
                     ..super::settings::Settings::default()
                 },
                 src: vec![test_resource_path("fixtures/isort")],

crates/ruff/src/rules/isort/rules/organize_imports.rs

@@ -122,14 +122,11 @@ pub fn organize_imports(
         &settings.src,
         package,
         settings.isort.combine_as_imports,
-        &settings.isort.extra_standard_library,
         settings.isort.force_single_line,
         settings.isort.force_sort_within_sections,
         settings.isort.force_wrap_aliases,
         &settings.isort.force_to_top,
-        &settings.isort.known_first_party,
+        &settings.isort.known_modules,
-        &settings.isort.known_third_party,
-        &settings.isort.known_local_folder,
         settings.isort.order_by_type,
         settings.isort.relative_imports_order,
         &settings.isort.single_line_exclusions,

crates/ruff/src/rules/isort/settings.rs

@@ -5,6 +5,7 @@ use std::collections::BTreeSet;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
+use crate::rules::isort::categorize::KnownModules;
 use ruff_macros::{CacheKey, ConfigurationOptions};
 
 use super::categorize::ImportType;
@@ -271,14 +272,11 @@ pub struct Options {
 pub struct Settings {
     pub required_imports: BTreeSet<String>,
     pub combine_as_imports: bool,
-    pub extra_standard_library: BTreeSet<String>,
     pub force_single_line: bool,
     pub force_sort_within_sections: bool,
     pub force_wrap_aliases: bool,
     pub force_to_top: BTreeSet<String>,
-    pub known_first_party: BTreeSet<String>,
+    pub known_modules: KnownModules,
-    pub known_third_party: BTreeSet<String>,
-    pub known_local_folder: BTreeSet<String>,
     pub order_by_type: bool,
     pub relative_imports_order: RelativeImportsOrder,
     pub single_line_exclusions: BTreeSet<String>,
@@ -297,14 +295,11 @@ impl Default for Settings {
         Self {
             required_imports: BTreeSet::new(),
             combine_as_imports: false,
-            extra_standard_library: BTreeSet::new(),
             force_single_line: false,
             force_sort_within_sections: false,
             force_wrap_aliases: false,
             force_to_top: BTreeSet::new(),
-            known_first_party: BTreeSet::new(),
+            known_modules: KnownModules::default(),
-            known_third_party: BTreeSet::new(),
-            known_local_folder: BTreeSet::new(),
             order_by_type: true,
             relative_imports_order: RelativeImportsOrder::default(),
             single_line_exclusions: BTreeSet::new(),
@@ -325,16 +320,16 @@ impl From<Options> for Settings {
         Self {
             required_imports: BTreeSet::from_iter(options.required_imports.unwrap_or_default()),
             combine_as_imports: options.combine_as_imports.unwrap_or(false),
-            extra_standard_library: BTreeSet::from_iter(
-                options.extra_standard_library.unwrap_or_default(),
-            ),
             force_single_line: options.force_single_line.unwrap_or(false),
             force_sort_within_sections: options.force_sort_within_sections.unwrap_or(false),
             force_wrap_aliases: options.force_wrap_aliases.unwrap_or(false),
             force_to_top: BTreeSet::from_iter(options.force_to_top.unwrap_or_default()),
-            known_first_party: BTreeSet::from_iter(options.known_first_party.unwrap_or_default()),
+            known_modules: KnownModules::new(
-            known_third_party: BTreeSet::from_iter(options.known_third_party.unwrap_or_default()),
+                options.known_first_party.unwrap_or_default(),
-            known_local_folder: BTreeSet::from_iter(options.known_local_folder.unwrap_or_default()),
+                options.known_third_party.unwrap_or_default(),
+                options.known_local_folder.unwrap_or_default(),
+                options.extra_standard_library.unwrap_or_default(),
+            ),
             order_by_type: options.order_by_type.unwrap_or(true),
             relative_imports_order: options.relative_imports_order.unwrap_or_default(),
             single_line_exclusions: BTreeSet::from_iter(
@@ -357,14 +352,20 @@ impl From<Settings> for Options {
         Self {
             required_imports: Some(settings.required_imports.into_iter().collect()),
             combine_as_imports: Some(settings.combine_as_imports),
-            extra_standard_library: Some(settings.extra_standard_library.into_iter().collect()),
+            extra_standard_library: Some(
+                settings
+                    .known_modules
+                    .standard_library
+                    .into_iter()
+                    .collect(),
+            ),
             force_single_line: Some(settings.force_single_line),
             force_sort_within_sections: Some(settings.force_sort_within_sections),
             force_wrap_aliases: Some(settings.force_wrap_aliases),
             force_to_top: Some(settings.force_to_top.into_iter().collect()),
-            known_first_party: Some(settings.known_first_party.into_iter().collect()),
+            known_first_party: Some(settings.known_modules.first_party.into_iter().collect()),
-            known_third_party: Some(settings.known_third_party.into_iter().collect()),
+            known_third_party: Some(settings.known_modules.third_party.into_iter().collect()),
-            known_local_folder: Some(settings.known_local_folder.into_iter().collect()),
+            known_local_folder: Some(settings.known_modules.local_folder.into_iter().collect()),
             order_by_type: Some(settings.order_by_type),
             relative_imports_order: Some(settings.relative_imports_order),
             single_line_exclusions: Some(settings.single_line_exclusions.into_iter().collect()),

crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__1_separate_subpackage_first_and_third_party_imports.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: 9
+    column: 0
+  fix:
+    edits:
+      - content: "import sys\n\nimport foo\nfrom foo import bar, baz\n\nimport baz\nimport foo.bar\nimport foo.bar.baz\nfrom foo.bar import blah, blub\nfrom foo.bar.baz import something\n"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 9
+          column: 0
+  parent: ~
+

crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__2_separate_subpackage_first_and_third_party_imports.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: 9
+    column: 0
+  fix:
+    edits:
+      - content: "import sys\n\nimport baz\nimport foo.bar\nimport foo.bar.baz\nfrom foo.bar import blah, blub\nfrom foo.bar.baz import something\n\nimport foo\nfrom foo import bar, baz\n"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 9
+          column: 0
+  parent: ~
+