From d51f34ee20342d26cb0961d9ded3ef77d02700c1 Mon Sep 17 00:00:00 2001
From: Ezra Shaw <ezrashawdev@gmail.com>
Date: Tue, 21 Nov 2023 21:43:02 +1300
Subject: [PATCH] fix: mark `__main__` as first-party import

---
 .../test/fixtures/isort/main_first_party.py   | 11 +++++++
 .../ruff_linter/src/rules/isort/categorize.rs |  5 +++
 crates/ruff_linter/src/rules/isort/mod.rs     | 24 ++++++++++++++
 ...__tests__sections_main_first_party.py.snap | 31 +++++++++++++++++++
 4 files changed, 71 insertions(+)
 create mode 100644 crates/ruff_linter/resources/test/fixtures/isort/main_first_party.py
 create mode 100644 crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__sections_main_first_party.py.snap

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
+
+