Merge pull request #350 from maggiemoss/add-config-option-to-collapse-lines

Add config option to collapse lines in same category

Author: maggiemoss

usort/config.py

@@ -68,6 +68,11 @@ class Config:
     # Whether to preserve inline comments on individual imports when sorting
     preserve_inline_comments: bool = False
 
+    # Whether to collapse blank lines within a single import category
+    # Default True matches current behavior (collapse blank lines within categories)
+    # Set to False to preserve one blank line within categories (pre-commit 58c01556 behavior)
+    collapse_blank_lines_in_category: bool = True
+
     # gitignore-style filename patterns to exclude when sorting entire directories
     excludes: List[str] = field(default_factory=list)
 
@@ -172,6 +177,10 @@ def update_from_config(self, toml_path: Path) -> None:
             self.merge_imports = bool(tbl["merge_imports"])
         if "preserve_inline_comments" in tbl:
             self.preserve_inline_comments = bool(tbl["preserve_inline_comments"])
+        if "collapse_blank_lines_in_category" in tbl:
+            self.collapse_blank_lines_in_category = bool(
+                tbl["collapse_blank_lines_in_category"]
+            )
         if "excludes" in tbl:
             self.excludes = tbl["excludes"]
 

usort/sorting.py

@@ -238,7 +238,11 @@ def fixup_whitespace(
             elif imp.sort_key.category_index != cur_category:
                 blanks = ("",)
             else:
-                blanks = ()
+                # Apply the config option for blank line behavior within categories
+                if self.config.collapse_blank_lines_in_category:
+                    blanks = ()
+                else:
+                    blanks = _old_blanks[:1]
 
             imp.comments.before = [*blanks, *old_comments]
 

usort/tests/functional.py

@@ -11,7 +11,7 @@
 from typing import Optional
 
 from ..api import usort, usort_path
-from ..config import Config
+from ..config import CAT_FIRST_PARTY, Config
 from ..translate import import_from_node
 from ..util import parse_import
 
@@ -1204,6 +1204,124 @@ def test_preserve_inline_comments_disabled(self) -> None:
             config,
         )
 
+    def test_collapse_blank_lines_in_category_enabled(self) -> None:
+        """Test with collapse_blank_lines_in_category=True (default, post-commit 58c01556 behavior)"""
+        config = replace(DEFAULT_CONFIG, collapse_blank_lines_in_category=True)
+        self.assertUsortResult(
+            """
+                import math
+
+                import gamma
+
+                import alpha
+
+
+                # special
+                import beta
+
+                from . import foo
+
+                import zeta
+
+            """,
+            """
+                import math
+
+                import alpha
+                # special
+                import beta
+                import gamma
+                import zeta
+
+                from . import foo
+
+            """,
+            config,
+        )
+
+    def test_collapse_blank_lines_in_category_disabled(self) -> None:
+        """Test with collapse_blank_lines_in_category=False (pre-commit 58c01556 behavior)"""
+        config = replace(DEFAULT_CONFIG, collapse_blank_lines_in_category=False)
+        self.assertUsortResult(
+            """
+                import math
+
+                import gamma
+
+                import alpha
+
+
+                # special
+                import beta
+
+                from . import foo
+
+                import zeta
+
+            """,
+            """
+                import math
+
+                import alpha
+
+                # special
+                import beta
+
+                import gamma
+
+                import zeta
+
+                from . import foo
+
+            """,
+            config,
+        )
+
+    def test_collapse_blank_lines_in_category_default(self) -> None:
+        # Configure torch as first_party so both imports are in the same category
+        known = DEFAULT_CONFIG.known.copy()
+        known["torch"] = CAT_FIRST_PARTY
+        config = replace(
+            DEFAULT_CONFIG,
+            collapse_blank_lines_in_category=True,
+            known=known,
+        )
+        self.assertUsortResult(
+            """
+                import torch
+
+                from .runner import get_nn_runners
+            """,
+            """
+                import torch
+                from .runner import get_nn_runners
+            """,
+            config,
+        )
+
+    def test_collapse_blank_lines_in_category_false(self) -> None:
+        # Configure torch as first_party so both imports are in the same category
+        known = DEFAULT_CONFIG.known.copy()
+        known["torch"] = CAT_FIRST_PARTY
+        config = replace(
+            DEFAULT_CONFIG,
+            collapse_blank_lines_in_category=False,
+            known=known,
+        )
+        self.assertUsortResult(
+            """
+                import torch
+
+                from .runner import get_runners
+            """,
+            """
+                import torch
+
+                from .runner import get_runners
+            """,
+            config,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()