fix: keep comments with statement in IndentNestedKeywords

docs/formatter/formatters/IndentNestedKeywords.md

@@ -92,6 +92,33 @@ You can configure it using ``indent_and``:
     ```
 
 
+## Comments
+
+Comments are kept together with the code they refer to. A comment that trails data on a line (for example a
+``# robocop: fmt: off`` disabler) stays on the same output line as the closest keyword or argument, instead of being
+moved in front of the whole statement:
+
+=== "Before"
+
+    ```robotframework
+    *** Test Cases ***
+    Test
+        Run Keywords    Log    1    AND    Log    2    # keep me next to log 2
+    ```
+
+=== "After"
+
+    ```robotframework
+    *** Test Cases ***
+    Test
+        Run Keywords
+        ...    Log    1
+        ...    AND
+        ...    Log    2    # keep me next to log 2
+    ```
+
+Comments that occupy their own line are kept as standalone comment lines placed before the formatted statement.
+
 ## Skip formatting settings
 
 To skip formatting run keywords inside settings (such as ``Suite Setup``, ``[Setup]``, ``[Teardown]`` etc.) set

src/robocop/formatter/formatters/IndentNestedKeywords.py

@@ -41,6 +41,10 @@ class IndentNestedKeywords(Formatter):
     ``AND`` argument inside ``Run Keywords`` can be handled in different ways. It is controlled via ``indent_and``
     parameter. For more details see the full documentation.
 
+    Comments are kept together with the code they reference: a comment trailing data on a line stays on the same
+    output line as the closest keyword or argument (rather than being moved in front of the whole statement), while a
+    comment occupying its own line is kept as a standalone comment line before the statement.
+
     To skip formatting run keywords inside settings (such as ``Suite Setup``, ``[Setup]``, ``[Teardown]`` etc.) set
     ``skip_settings`` to ``True``.
     """
@@ -81,39 +85,63 @@ def get_separator(self, column: int = 1, continuation: bool = False) -> Token:
             separator = self.formatting_config.separator * column
         return Token(Token.SEPARATOR, separator)
 
+    def append_trailing_comments(
+        self, tokens: list[Token], line: list[Token], anchor_map: dict[int, list[Token]]
+    ) -> None:
+        """Append comments anchored to any data token in ``line`` at the end of the current output line."""
+        if not anchor_map:
+            return
+        for data_token in line:
+            for comment in anchor_map.get(id(data_token), []):
+                tokens.append(self.get_separator())
+                tokens.append(comment)
+
     def parse_keyword_lines(
-        self, lines: list[tuple[int, list[Token]]], tokens: list[Token], new_line: list[Token], eol: Token
+        self,
+        lines: list[tuple[int, list[Token]]],
+        tokens: list[Token],
+        new_line: list[Token],
+        eol: Token,
+        anchor_map: dict[int, list[Token]] | None = None,
     ) -> list[Token]:
+        anchor_map = anchor_map or {}
         separator = self.get_separator()
+        self.append_trailing_comments(tokens, lines[0][1], anchor_map)
         for column, line in lines[1:]:
             tokens.extend(new_line)
             tokens.append(self.get_separator(column, continuation=True))
             tokens.extend(misc.join_tokens_with_token(line, separator))
+            self.append_trailing_comments(tokens, line, anchor_map)
         tokens.append(eol)
         return tokens
 
     @staticmethod
-    def node_was_formatted(old_tokens: list[Token], new_tokens: list[Token]) -> bool:
-        """Compare code before and after formatting while ignoring comments to check if code was formatted."""
-        if len(new_tokens) > len(old_tokens):
-            return True
-        old_tokens_no_comm: list[Token] = []
+    def tokens_without_comments(tokens: list[Token]) -> list[Token]:
+        """Return tokens with comments (and resulting comment-only lines) removed."""
+        result: list[Token] = []
         data_in_line = False
-        for token in old_tokens:
+        for token in tokens:
             if token.type == Token.EOL:
                 if not data_in_line:
                     continue
                 data_in_line = False
             elif token.type == Token.COMMENT:
-                if old_tokens_no_comm and old_tokens_no_comm[-1].type == Token.SEPARATOR:
-                    old_tokens_no_comm.pop()
+                if result and result[-1].type == Token.SEPARATOR:
+                    result.pop()
                 continue
             elif token.type != Token.SEPARATOR:
                 data_in_line = True
-            old_tokens_no_comm.append(token)
-        if len(new_tokens) != len(old_tokens_no_comm):
+            result.append(token)
+        return result
+
+    @classmethod
+    def node_was_formatted(cls, old_tokens: list[Token], new_tokens: list[Token]) -> bool:
+        """Compare code before and after formatting while ignoring comments to check if code was formatted."""
+        old_tokens_no_comm = cls.tokens_without_comments(old_tokens)
+        new_tokens_no_comm = cls.tokens_without_comments(new_tokens)
+        if len(new_tokens_no_comm) != len(old_tokens_no_comm):
             return True
-        for new_token, old_token in zip(new_tokens, old_tokens_no_comm, strict=False):
+        for new_token, old_token in zip(new_tokens_no_comm, old_tokens_no_comm, strict=False):
             if new_token.type != old_token.type or new_token.value != old_token.value:
                 return True
         return False
@@ -123,11 +151,11 @@ def visit_SuiteSetup(self, node: SuiteSetup) -> SuiteSetup | tuple[Any, ...]:  #
         lines = self.get_setting_lines(node, 0)
         if not lines:
             return node
-        comments = misc.collect_comments_from_tokens(node.tokens, indent=None)
+        comments, anchor_map = misc.split_comments_by_anchor(node.tokens, indent=None)
         separator = self.get_separator()
         new_line = misc.get_new_line()
         tokens = [node.data_tokens[0], separator, *misc.join_tokens_with_token(lines[0][1], separator)]
-        formatted_tokens = self.parse_keyword_lines(lines, tokens, new_line, eol=node.tokens[-1])
+        formatted_tokens = self.parse_keyword_lines(lines, tokens, new_line, eol=node.tokens[-1], anchor_map=anchor_map)
         if self.node_was_formatted(node.tokens, formatted_tokens):
             node.tokens = formatted_tokens
             return (*comments, node)
@@ -136,7 +164,7 @@ def visit_SuiteSetup(self, node: SuiteSetup) -> SuiteSetup | tuple[Any, ...]:  #
     visit_SuiteTeardown = visit_TestSetup = visit_TestTeardown = visit_SuiteSetup  # noqa: N815
 
     @skip_if_disabled
-    def visit_Setup(self, node: Setup) -> Setup:  # noqa: N802
+    def visit_Setup(self, node: Setup) -> Setup | tuple[Any, ...]:  # noqa: N802
         indent = len(node.tokens[0].value)
         lines = self.get_setting_lines(node, indent)
         if not lines:
@@ -145,12 +173,10 @@ def visit_Setup(self, node: Setup) -> Setup:  # noqa: N802
         separator = self.get_separator()
         new_line = misc.get_new_line(indent)
         tokens = [indent, node.data_tokens[0], separator, *misc.join_tokens_with_token(lines[0][1], separator)]
-        comment = misc.merge_comments_into_one(node.tokens)
-        if comment:
-            # need to add comments on first line for [Setup] / [Teardown] settings
-            comment_sep = Token(Token.SEPARATOR, "  ")
-            tokens.extend([comment_sep, comment])
-        node.tokens = self.parse_keyword_lines(lines, tokens, new_line, eol=node.tokens[-1])
+        comments, anchor_map = misc.split_comments_by_anchor(node.tokens, indent=indent)
+        node.tokens = self.parse_keyword_lines(lines, tokens, new_line, eol=node.tokens[-1], anchor_map=anchor_map)
+        if comments:
+            return (*comments, node)
         return node
 
     visit_Teardown = visit_Setup  # noqa: N815
@@ -164,7 +190,7 @@ def visit_KeywordCall(self, node: KeywordCall) -> KeywordCall | tuple[Any, ...]:
             return node
 
         indent = node.tokens[0]
-        comments = misc.collect_comments_from_tokens(node.tokens, indent)
+        comments, anchor_map = misc.split_comments_by_anchor(node.tokens, indent)
         assign, kw_tokens = misc.split_on_token_type(node.data_tokens, Token.KEYWORD)
         lines = self.parse_sub_kw(kw_tokens)
         if not lines:
@@ -177,7 +203,7 @@ def visit_KeywordCall(self, node: KeywordCall) -> KeywordCall | tuple[Any, ...]:
             tokens.extend([*misc.join_tokens_with_token(assign, separator), separator])
         tokens.extend(misc.join_tokens_with_token(lines[0][1], separator))
         new_line = misc.get_new_line(indent)
-        formatted_tokens = self.parse_keyword_lines(lines, tokens, new_line, eol=node.tokens[-1])
+        formatted_tokens = self.parse_keyword_lines(lines, tokens, new_line, eol=node.tokens[-1], anchor_map=anchor_map)
         if self.node_was_formatted(node.tokens, formatted_tokens):
             node.tokens = formatted_tokens
             return (*comments, node)

src/robocop/formatter/utils/misc.py

@@ -371,6 +371,49 @@ def collect_comments_from_tokens(tokens: list[Token], indent: Token | None) -> l
     return [Comment([comment, eol]) for comment in comments]
 
 
+def _split_tokens_into_lines(tokens: list[Token]) -> Generator[list[Token]]:
+    """Yield tokens grouped by physical line (split on EOL)."""
+    line: list[Token] = []
+    for token in tokens:
+        line.append(token)
+        if token.type == Token.EOL:
+            yield line
+            line = []
+    if line:
+        yield line
+
+
+def split_comments_by_anchor(tokens: list[Token], indent: Token | None) -> tuple[list[Comment], dict[int, list[Token]]]:
+    """
+    Split comments into standalone comment lines and trailing comments.
+
+    A comment that shares a physical line with data tokens (keyword name or arguments) is considered a
+    *trailing* comment and is anchored to the last data token preceding it on that line. A comment that occupies
+    its own line is considered *standalone*.
+
+    Returns a tuple of ``(standalone_comments, anchor_map)`` where ``standalone_comments`` is a list of
+    ``Comment`` nodes (to be emitted before the statement) and ``anchor_map`` maps ``id(data_token)`` to the list
+    of trailing comment tokens that should be rendered on the same output line as that data token.
+    """
+    standalone: list[Comment] = []
+    anchor_map: dict[int, list[Token]] = {}
+    eol = Token(Token.EOL)
+    for line in _split_tokens_into_lines(tokens):
+        data_tokens = [
+            token for token in line if token.type not in (Token.SEPARATOR, Token.EOL, Token.CONTINUATION, Token.COMMENT)
+        ]
+        line_comments = get_comments(line)
+        if not line_comments:
+            continue
+        if data_tokens:
+            anchor_map.setdefault(id(data_tokens[-1]), []).extend(line_comments)
+        elif indent:
+            standalone.extend(Comment([indent, comment, eol]) for comment in line_comments)
+        else:
+            standalone.extend(Comment([comment, eol]) for comment in line_comments)
+    return standalone, anchor_map
+
+
 def flatten_multiline(tokens: list[Token], separator: str, remove_comments: bool = False) -> list[Token]:
     flattened = []
     skip_start = False

tests/formatter/formatters/IndentNestedKeywords/expected/comments.robot

@@ -2,28 +2,21 @@
 Suite Setup    Run Keywords
 ...    No Operation
 # ...    No Operation
-# comment
 Suite Teardown    Run Keywords
 ...    No Operation
-...    No Operation
+...    No Operation    # comment
 Test Setup    Run Keywords
 # ...    No Operation
 ...    No Operation
 
 
 *** Keywords ***
 Comments
-    # comment1  comment2
-    # comment 3
-    # comment 4
-    # comment 5
-    # comment 6
-    # comment 7
-    Run Keyword
-    ...    Run Keyword If    ${True}
-    ...        Keyword    ${arg}
-    ...    ELSE
-    ...        Keyword
+    Run Keyword    # comment1  comment2
+    ...    Run Keyword If    ${True}    # comment 3
+    ...        Keyword    ${arg}    # comment 4    # comment 5
+    ...    ELSE    # comment 6
+    ...        Keyword    # comment 7
 
 Golden Keywords
     [Test Setup]    Run Keywords
@@ -39,3 +32,9 @@ Golden Keywords
     Run Keywords
     ...    No Operation
     ...    No Operation  # comment about this keyword call
+
+Trailing Comment Stays With Statement
+    Run Keywords
+    ...    Log    1
+    ...    AND
+    ...    Log    2    # keep me next to log 2

tests/formatter/formatters/IndentNestedKeywords/expected/settings.robot

@@ -2,31 +2,29 @@
 Suite Setup    Run Keywords
 ...    No Operation
 ...    No Operation
-# comment1
 Suite Teardown    Run Keywords
 ...    Log    1
 ...    AND
-...    Log    2
+...    Log    2    # comment1
 
 Test Setup    Run Keywords
 ...    No Operation
 ...    No Operation
-# comment1  comment2
 Test Teardown    Run Keywords
 ...    No Operation
-...    No Operation
+...    No Operation    # comment1  comment2
 
 
 *** Test Cases ***
 Test
     [Setup]    Run Keyword If    ${True}
     ...    No Operation
-    [Teardown]    Run Keywords  # comment comment2
-    ...    Log    1
+    [Teardown]    Run Keywords
+    ...    Log    1    # comment
     ...    AND
-    ...    Log    2
+    ...    Log    2    # comment2
     No Operation
 
-Test in line    [Setup]    Run Keyword  # comment
-    ...    Log    1
+Test in line    [Setup]    Run Keyword
+    ...    Log    1    # comment
     No Operation

tests/formatter/formatters/IndentNestedKeywords/expected/too_long_line.robot

@@ -2,20 +2,18 @@
 Suite Setup    Run Keywords
 ...    No Operation
 ...    No Operation
-# comment1
 Suite Teardown    Run Keywords
 ...    Log    1
 ...    AND
 ...    Keyword Call That Is Going Over Allowed Limit    ${myVeryLongDefinitionOfAnElement}
-...        ${myVeryLongDefinitionOfAnElement}
+...        ${myVeryLongDefinitionOfAnElement}    # comment1
 
 Test Setup    Run Keywords
 ...    No Operation
 ...    Keyword Call That Is Going Over Allowed Limit Keyword Call That Is Going Over Allowed Limit Keyword Call That Is Going Over Allowed Limit
-# comment1  comment2
 Test Teardown    Run Keywords
 ...    Keyword Call That Is Going Over Allowed Limit Keyword Call That Is Going Over Allowed Limit Keyword Call That Is Going Over Allowed Limit
-...    No Operation
+...    No Operation    # comment1  comment2
 
 
 *** Test Cases ***
@@ -77,7 +75,7 @@ Keyword That Should Be Split On Every Line
     ...            ${addingTheSecondArgumentMakesThisLineTooLong}
 
 Settings
-    [Teardown]    Run Keywords  # comment comment2
+    [Teardown]    Run Keywords    # comment comment2
     ...    Element Should Contain    ${myVeryLongDefinitionOfAnElement}
     ...        ${addingTheSecondArgumentMakesThisLineTooLong}
     ...        ${addingTheSecondArgumentMakesThisLineTooLong}

tests/formatter/formatters/IndentNestedKeywords/source/comments.robot

@@ -31,3 +31,6 @@ Golden Keywords
     Run Keywords
     ...    No Operation
     ...    No Operation  # comment about this keyword call
+
+Trailing Comment Stays With Statement
+    Run Keywords    Log    1    AND    Log    2    # keep me next to log 2