diff --git a/crates/ruff_linter/resources/test/fixtures/isort/main_first_party.py b/crates/ruff_linter/resources/test/fixtures/isort/main_first_party.py new file mode 100644 index 0000000000000..ba224527c3b9a --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/isort/main_first_party.py @@ -0,0 +1,11 @@ +import os + +import __main__ +import third_party + +import first_party + +os.a +third_party.a +__main__.a +first_party.a diff --git a/crates/ruff_linter/src/rules/isort/categorize.rs b/crates/ruff_linter/src/rules/isort/categorize.rs index 6c086bfb2f7e2..8ab090645ab62 100644 --- a/crates/ruff_linter/src/rules/isort/categorize.rs +++ b/crates/ruff_linter/src/rules/isort/categorize.rs @@ -106,6 +106,11 @@ pub(crate) fn categorize<'a>( &ImportSection::Known(ImportType::FirstParty), Reason::SourceMatch(src), ) + } else if matches!(level, None | Some(0)) && module_name == "__main__" { + ( + &ImportSection::Known(ImportType::FirstParty), + Reason::KnownFirstParty, + ) } else { ( &ImportSection::Known(ImportType::ThirdParty), diff --git a/crates/ruff_linter/src/rules/isort/mod.rs b/crates/ruff_linter/src/rules/isort/mod.rs index 73380ca4eefd0..10a3f7b1d974a 100644 --- a/crates/ruff_linter/src/rules/isort/mod.rs +++ b/crates/ruff_linter/src/rules/isort/mod.rs @@ -1033,6 +1033,30 @@ mod tests { Ok(()) } + #[test_case(Path::new("main_first_party.py"))] + fn main_is_first_party(path: &Path) -> Result<()> { + let snapshot = format!("sections_{}", path.to_string_lossy()); + let diagnostics = test_path( + Path::new("isort").join(path).as_path(), + &LinterSettings { + src: vec![test_resource_path("fixtures/isort")], + isort: super::settings::Settings { + known_modules: KnownModules::new( + vec![pattern("first_party")], + vec![], + vec![], + vec![], + FxHashMap::default(), + ), + ..super::settings::Settings::default() + }, + ..LinterSettings::for_rule(Rule::UnsortedImports) + }, + )?; + assert_messages!(snapshot, diagnostics); + Ok(()) + } + #[test] fn detect_same_package() -> Result<()> { let diagnostics = test_path( diff --git a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__sections_main_first_party.py.snap b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__sections_main_first_party.py.snap new file mode 100644 index 0000000000000..404ad962cb670 --- /dev/null +++ b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__sections_main_first_party.py.snap @@ -0,0 +1,31 @@ +--- +source: crates/ruff_linter/src/rules/isort/mod.rs +--- +main_first_party.py:1:1: I001 [*] Import block is un-sorted or un-formatted + | + 1 | / import os + 2 | | + 3 | | import __main__ + 4 | | import third_party + 5 | | + 6 | | import first_party + 7 | | + 8 | | os.a + | |_^ I001 + 9 | third_party.a +10 | __main__.a + | + = help: Organize imports + +ℹ Safe fix +1 1 | import os +2 2 | +3 |-import __main__ +4 3 | import third_party +5 4 | + 5 |+import __main__ +6 6 | import first_party +7 7 | +8 8 | os.a + +