Merge pull request #344 from maggiemoss/preserve-inline-comments

Preserve inline comments

Author: maggiemoss

docs/guide.rst

@@ -537,6 +537,57 @@ The following options are valid for the main ``tool.usort`` table:
     Whether to merge sequential imports from the same base module.
     See `Merging`_ for details on how this works.
 
+.. attribute:: preserve_inline_comments
+    :type: bool
+    :value: false
+
+    Whether to preserve the placement of inline comments with their associated
+    import items when sorting and merging imports.
+
+    When set to ``false`` (the default), inline comments from single-line imports
+    are treated as import-level comments and will be placed at the end of the
+    merged import statement::
+
+        # Before
+        from foo import bar  # comment1
+        from foo import baz  # comment2
+
+        # After (default behavior)
+        from foo import bar, baz  # comment1  # comment2
+
+    When set to ``true``, inline comments are associated with their specific
+    import items and will remain attached to those items after sorting and
+    merging::
+
+        # Before
+        from foo import bar  # comment1
+        from foo import baz  # comment2
+
+        # After (preserve_inline_comments = true)
+        from foo import (
+            bar,  # comment1
+            baz,  # comment2
+        )
+
+    This is particularly useful when using inline comments for directives like
+    type checker annotations or linter suppressions that are specific to individual
+    imports::
+
+        # Before
+        from torch._C import _linalg  # pyrefly:
+        from torch._C import _LinAlgError as LinAlgError  # pyrefly: ignore
+        from torch._C import _add_docstr
+
+        # After (preserve_inline_comments = true)
+        from torch._C import (
+            _add_docstr,
+            _linalg,  # pyrefly:
+            _LinAlgError as LinAlgError,  # pyrefly: ignore
+        )
+
+    Note that this option only affects inline comments on single-line imports.
+    Comments on multi-line imports are always preserved with their associated items.
+
 .. attribute:: excludes
     :type: List[str]
 

usort/api.py

@@ -16,7 +16,6 @@
 from .types import Result
 from .util import get_timings, timed, try_parse
 
-
 __all__ = ["usort_bytes", "usort_string", "usort_path", "usort_stdin"]
 
 

usort/config.py

@@ -65,6 +65,9 @@ class Config:
     # Whether to merge imports when sorting
     merge_imports: bool = True
 
+    # Whether to preserve inline comments on individual imports when sorting
+    preserve_inline_comments: bool = False
+
     # gitignore-style filename patterns to exclude when sorting entire directories
     excludes: List[str] = field(default_factory=list)
 
@@ -167,6 +170,8 @@ def update_from_config(self, toml_path: Path) -> None:
             self.magic_commas = bool(tbl["magic_commas"])
         if "merge_imports" in tbl:
             self.merge_imports = bool(tbl["merge_imports"])
+        if "preserve_inline_comments" in tbl:
+            self.preserve_inline_comments = bool(tbl["preserve_inline_comments"])
         if "excludes" in tbl:
             self.excludes = tbl["excludes"]
 

usort/tests/cli.py

@@ -332,8 +332,8 @@ def test_format_latin_1(self) -> None:
 s = "\xb5"
 """.replace(
                 b"\r", b""
-            )  # git on windows might make \r\n
-        ) as dtmp:
+            )
+        ) as dtmp:  # git on windows might make \r\n
             runner = CliRunner()
 
             # Diff output is unicode

usort/tests/functional.py

@@ -1140,6 +1140,58 @@ def test_sorting_with_extra_blank_lines(self) -> None:
             """,
         )
 
+    def test_preserve_inline_comments(self) -> None:
+        config = replace(DEFAULT_CONFIG, preserve_inline_comments=True)
+        self.assertUsortResult(
+            """
+                from torch._C import _linalg  # pyrefly:
+                from torch._C import (
+                    _LinAlgError as LinAlgError,  # pyrefly: ignore  # missing-module-attribute
+                )
+                from torch._C import _add_docstr
+            """,
+            """
+                from torch._C import (
+                    _add_docstr,
+                    _linalg,  # pyrefly:
+                    _LinAlgError as LinAlgError,  # pyrefly: ignore  # missing-module-attribute
+                )
+            """,
+            config,
+        )
+
+    def test_preserve_inline_comments_multiple_modules(self) -> None:
+        config = replace(DEFAULT_CONFIG, preserve_inline_comments=True)
+        self.assertUsortResult(
+            """
+                from foo import bar  # comment1
+                from foo import baz  # comment2
+                from foo import alpha
+            """,
+            """
+                from foo import (
+                    alpha,
+                    bar,  # comment1
+                    baz,  # comment2
+                )
+            """,
+            config,
+        )
+
+    def test_preserve_inline_comments_disabled(self) -> None:
+        # Without preserve_inline_comments, comments become import-level
+        config = replace(DEFAULT_CONFIG, preserve_inline_comments=False)
+        self.assertUsortResult(
+            """
+                from torch._C import _linalg  # pyrefly:
+                from torch._C import _add_docstr
+            """,
+            """
+                from torch._C import _add_docstr, _linalg  # pyrefly:
+            """,
+            config,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

usort/translate.py

@@ -201,10 +201,27 @@ def import_from_node(node: cst.SimpleStatementLine, config: Config) -> SortableI
         if node.body[0].relative:
             stem = "." * len(node.body[0].relative) + stem
 
-        for name in node.body[0].names:
+        # For single-line imports with preserve_inline_comments, attach trailing
+        # comments to the specific item instead of the import level
+        is_single_line = not node.body[0].lpar
+        single_item_count = len(node.body[0].names)
+
+        for idx, name in enumerate(node.body[0].names):
             items.append(item_from_node(name, stem, comments.initial))
             comments.initial = []
 
+            # If preserving inline comments and this is a single-line import with
+            # one item, move first_inline comments to the item's inline comments
+            if (
+                config.preserve_inline_comments
+                and is_single_line
+                and single_item_count == 1
+                and idx == 0
+                and comments.first_inline
+            ):
+                items[-1].comments.inline.extend(comments.first_inline)
+                comments.first_inline.clear()
+
     else:
         raise TypeError
 
@@ -240,6 +257,14 @@ def import_to_node(
 ) -> cst.BaseStatement:
     if config.magic_commas and imp.stem and imp.trailing_comma:
         return import_to_node_multi(imp, module)
+
+    # If preserve_inline_comments is enabled and any item has inline comments,
+    # use multi-line format to preserve them
+    if config.preserve_inline_comments and imp.stem:
+        has_item_comments = any(item.comments.inline for item in imp.items)
+        if has_item_comments:
+            return import_to_node_multi(imp, module)
+
     node = import_to_node_single(imp, module)
     content = indent + render_node(node, module).rstrip()
     # basic imports can't be reflowed, so only deal with from-imports