feat(isort): Implement known-local-folder

Related to #2419
astral-sh · Feb 10, 2023 · b171dbf · b171dbf
1 parent cda2ff0
commit b171dbf
Show file tree

Hide file tree

Showing 10 changed files with 72 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -3467,6 +3467,25 @@ known-first-party = ["src"]
 
 ---
 
+#### [`known-local-folder`](#known-local-folder)
+
+A list of modules to consider being a local folder.
+Generally, this is reserved for relative
+imports (from . import module).
+
+**Default value**: `[]`
+
+**Type**: `list[str]`
+
+**Example usage**:
+
+```toml
+[tool.ruff.isort]
+known-local-folder = ["src"]
+```
+
+---
+
 #### [`known-third-party`](#known-third-party)
 
 A list of modules to consider third-party, regardless of whether they

diff --git a/crates/ruff/resources/test/fixtures/isort/pyproject.toml b/crates/ruff/resources/test/fixtures/isort/pyproject.toml
@@ -3,3 +3,4 @@ line-length = 88
 
 [tool.ruff.isort]
 lines-after-imports = 3
+known-local-folder = ["ruff"]
diff --git a/crates/ruff/resources/test/fixtures/isort/separate_local_folder_imports.py b/crates/ruff/resources/test/fixtures/isort/separate_local_folder_imports.py
@@ -1,4 +1,5 @@
 import sys
+import ruff
 import leading_prefix
 import os
 from . import leading_prefix
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
@@ -164,6 +164,7 @@ pub fn typing_only_runtime_import(
             package,
             &settings.isort.known_first_party,
             &settings.isort.known_third_party,
+            &settings.isort.known_local_folder,
             &settings.isort.extra_standard_library,
             settings.target_version,
         ) {

diff --git a/crates/ruff/src/rules/isort/categorize.rs b/crates/ruff/src/rules/isort/categorize.rs
@@ -27,6 +27,7 @@ enum Reason<'a> {
     NonZeroLevel,
     KnownFirstParty,
     KnownThirdParty,
+    KnownLocalFolder,
     ExtraStandardLibrary,
     Future,
     KnownStandardLibrary,
@@ -43,6 +44,7 @@ pub fn categorize(
     package: Option<&Path>,
     known_first_party: &BTreeSet<String>,
     known_third_party: &BTreeSet<String>,
+    known_local_folder: &BTreeSet<String>,
     extra_standard_library: &BTreeSet<String>,
     target_version: PythonVersion,
 ) -> ImportType {
@@ -53,6 +55,8 @@ pub fn categorize(
             (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__" {
@@ -98,12 +102,14 @@ fn match_sources<'a>(paths: &'a [PathBuf], base: &str) -> Option<&'a Path> {
     None
 }
 
+#[allow(clippy::too_many_arguments)]
 pub fn categorize_imports<'a>(
     block: ImportBlock<'a>,
     src: &[PathBuf],
     package: Option<&Path>,
     known_first_party: &BTreeSet<String>,
     known_third_party: &BTreeSet<String>,
+    known_local_folder: &BTreeSet<String>,
     extra_standard_library: &BTreeSet<String>,
     target_version: PythonVersion,
 ) -> BTreeMap<ImportType, ImportBlock<'a>> {
@@ -117,6 +123,7 @@ pub fn categorize_imports<'a>(
             package,
             known_first_party,
             known_third_party,
+            known_local_folder,
             extra_standard_library,
             target_version,
         );
@@ -135,6 +142,7 @@ pub fn categorize_imports<'a>(
             package,
             known_first_party,
             known_third_party,
+            known_local_folder,
             extra_standard_library,
             target_version,
         );
@@ -153,6 +161,7 @@ pub fn categorize_imports<'a>(
             package,
             known_first_party,
             known_third_party,
+            known_local_folder,
             extra_standard_library,
             target_version,
         );
@@ -171,6 +180,7 @@ pub fn categorize_imports<'a>(
             package,
             known_first_party,
             known_third_party,
+            known_local_folder,
             extra_standard_library,
             target_version,
         );

diff --git a/crates/ruff/src/rules/isort/mod.rs b/crates/ruff/src/rules/isort/mod.rs
@@ -122,6 +122,7 @@ pub fn format_imports(
     force_wrap_aliases: bool,
     known_first_party: &BTreeSet<String>,
     known_third_party: &BTreeSet<String>,
+    known_local_folder: &BTreeSet<String>,
     order_by_type: bool,
     relative_imports_order: RelativeImportsOrder,
     single_line_exclusions: &BTreeSet<String>,
@@ -155,6 +156,7 @@ pub fn format_imports(
             force_wrap_aliases,
             known_first_party,
             known_third_party,
+            known_local_folder,
             order_by_type,
             relative_imports_order,
             single_line_exclusions,
@@ -212,6 +214,7 @@ fn format_import_block(
     force_wrap_aliases: bool,
     known_first_party: &BTreeSet<String>,
     known_third_party: &BTreeSet<String>,
+    known_local_folder: &BTreeSet<String>,
     order_by_type: bool,
     relative_imports_order: RelativeImportsOrder,
     single_line_exclusions: &BTreeSet<String>,
@@ -229,6 +232,7 @@ fn format_import_block(
         package,
         known_first_party,
         known_third_party,
+        known_local_folder,
         extra_standard_library,
         target_version,
     );
@@ -366,6 +370,12 @@ mod tests {
             Path::new("isort").join(path).as_path(),
             &Settings {
                 src: vec![test_resource_path("fixtures/isort")],
+                isort: super::settings::Settings {
+                    known_local_folder: vec!["ruff".to_string()]
+                        .into_iter()
+                        .collect::<BTreeSet<_>>(),
+                    ..super::settings::Settings::default()
+                },
                 ..Settings::for_rule(Rule::UnsortedImports)
             },
         )?;

diff --git a/crates/ruff/src/rules/isort/rules/organize_imports.rs b/crates/ruff/src/rules/isort/rules/organize_imports.rs
@@ -127,6 +127,7 @@ pub fn organize_imports(
         settings.isort.force_wrap_aliases,
         &settings.isort.known_first_party,
         &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,

diff --git a/crates/ruff/src/rules/isort/settings.rs b/crates/ruff/src/rules/isort/settings.rs
@@ -138,6 +138,17 @@ pub struct Options {
     /// A list of modules to consider third-party, regardless of whether they
     /// can be identified as such via introspection of the local filesystem.
     pub known_third_party: Option<Vec<String>>,
+    #[option(
+        default = r#"[]"#,
+        value_type = "list[str]",
+        example = r#"
+            known-local-folder = ["src"]
+        "#
+    )]
+    /// A list of modules to consider being a local folder.
+    /// Generally, this is reserved for relative
+    /// imports (from . import module).
+    pub known_local_folder: Option<Vec<String>>,
     #[option(
         default = r#"[]"#,
         value_type = "list[str]",
@@ -247,6 +258,7 @@ pub struct Settings {
     pub force_wrap_aliases: bool,
     pub known_first_party: BTreeSet<String>,
     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>,
@@ -270,6 +282,7 @@ impl Default for Settings {
             force_wrap_aliases: false,
             known_first_party: BTreeSet::new(),
             known_third_party: BTreeSet::new(),
+            known_local_folder: BTreeSet::new(),
             order_by_type: true,
             relative_imports_order: RelativeImportsOrder::default(),
             single_line_exclusions: BTreeSet::new(),
@@ -297,6 +310,7 @@ impl From<Options> for Settings {
             force_wrap_aliases: options.force_wrap_aliases.unwrap_or(false),
             known_first_party: BTreeSet::from_iter(options.known_first_party.unwrap_or_default()),
             known_third_party: BTreeSet::from_iter(options.known_third_party.unwrap_or_default()),
+            known_local_folder: BTreeSet::from_iter(options.known_local_folder.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(
@@ -324,6 +338,7 @@ impl From<Settings> for Options {
             force_wrap_aliases: Some(settings.force_wrap_aliases),
             known_first_party: Some(settings.known_first_party.into_iter().collect()),
             known_third_party: Some(settings.known_third_party.into_iter().collect()),
+            known_local_folder: Some(settings.known_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()),

diff --git a/...rc/rules/isort/snapshots/ruff__rules__isort__tests__separate_local_folder_imports.py.snap b/...rc/rules/isort/snapshots/ruff__rules__isort__tests__separate_local_folder_imports.py.snap
@@ -1,5 +1,5 @@
 ---
-source: src/rules/isort/mod.rs
+source: crates/ruff/src/rules/isort/mod.rs
 expression: diagnostics
 ---
 - kind:
@@ -8,7 +8,7 @@ expression: diagnostics
     row: 1
     column: 0
   end_location:
-    row: 5
+    row: 6
     column: 0
   fix:
     content:
@@ -17,13 +17,14 @@ expression: diagnostics
       - ""
       - import leading_prefix
       - ""
+      - import ruff
       - from . import leading_prefix
       - ""
     location:
       row: 1
       column: 0
     end_location:
-      row: 5
+      row: 6
       column: 0
   parent: ~
 
diff --git a/ruff.schema.json b/ruff.schema.json
@@ -964,6 +964,16 @@
             "type": "string"
           }
         },
+        "known-local-folder": {
+          "description": "A list of modules to consider being a local folder. Generally, this is reserved for relative imports (from . import module).",
+          "type": [
+            "array",
+            "null"
+          ],
+          "items": {
+            "type": "string"
+          }
+        },
         "known-third-party": {
           "description": "A list of modules to consider third-party, regardless of whether they can be identified as such via introspection of the local filesystem.",
           "type": [