OrderSettings now preserves comments order and groups them with settings (#488)

* OrderSettings now preserver comments order and group them with settings

* Update docs/releasenotes/4.0.0.rst

Co-authored-by: Mateusz Nojek <matnojek@gmail.com>

Author: bhirsz

docs/releasenotes/4.0.0.rst

@@ -116,3 +116,36 @@ rather than the sections you want to format. For example following command::
 Is now equivalent of::
 
     > robotidy -c NormalizeSeparators:skip_sections=keywords
+
+Group comments with settings in OrderSettings #468
+---------------------------------------------------
+
+``OrderSettings`` transformer adjusts the order of the settings such as ``[Arguments]`` or ``[Teardown]`` inside test,
+task or keyword. Previously, it only ordered the settings and comments were not moved. Following code::
+
+    *** Keywords ***
+    Keyword
+        # comment about step
+        Step
+        # comment about arguments
+        [Arguments]    ${arg}
+
+would result in::
+
+
+    *** Keywords ***
+    Keyword
+        [Arguments]    ${arg}
+        # comment about step
+        Step
+         # comment about arguments
+
+However in most cases such comments are meant only the very next line. That's why now the comments will be moved
+together with settings::
+
+    *** Keywords ***
+    Keyword
+        # comment about arguments
+        [Arguments]    ${arg}
+        # comment about step
+        Step

docs/source/transformers/OrderSettings.rst

@@ -44,6 +44,9 @@ settings like ``[Teardown]``, ``[Return]`` are moved to the end of keyword.
 Test case settings ``[Documentation]``, ``[Tags]``, ``[Template]``, ``[Timeout]``, ``[Setup]`` are put before test case body and
 ``[Teardown]`` is moved to the end of test case.
 
+Configure order of the settings
+----------------------------------
+
 Default order can be changed using following parameters:
 
 - ``keyword_before = documentation,tags,timeout,arguments``
@@ -103,3 +106,32 @@ This configuration is invalid because teardown is by default part of the ``test_
 We need to overwrite both orders::
 
     robotidy --configure OrderSettings:test_before=teardown:test_after=
+
+Settings comments
+---------------------
+
+Comments next to settings will be moved together.
+
+.. tab-set::
+
+    .. tab-item:: Before
+
+        .. code:: robotframework
+
+            *** Keywords ***
+            Keyword
+                # comment about step
+                Step
+                # comment about arguments
+                [Arguments]    ${arg}
+
+    .. tab-item:: After
+
+        .. code:: robotframework
+
+            *** Keywords ***
+            Keyword
+                # comment about arguments
+                [Arguments]    ${arg}
+                # comment about step
+                Step

robotidy/transformers/OrderSettings.py

@@ -22,7 +22,10 @@ def __init__(self, transformer, param_name, param_value):
 class SettingInBothOrdersError(RobotidyConfigError):
     def __init__(self, transformer, first_order, second_order, duplicates):
         names = ",".join(setting.lower() for setting in duplicates)
-        msg = f"{transformer}: Invalid '{first_order}' and '{second_order}' order values. Following setting names exists in both orders: {names}"
+        msg = (
+            f"{transformer}: Invalid '{first_order}' and '{second_order}' order values. "
+            f"Following setting names exists in both orders: {names}"
+        )
         super().__init__(msg)
 
 
@@ -156,24 +159,33 @@ def order_settings(self, node, setting_types, before, after):
         after_seen = False
         # when after_seen is set to True then all statements go to trailing_after and last non data
         # will be appended after tokens defined in `after` set (like [Return])
-        comment = []
+        comments, header_line = [], []
         for child in node.body:
-            if isinstance(child, Comment) and child.lineno == node.lineno:
-                comment.append(child)
+            if isinstance(child, Comment):
+                if child.lineno == node.lineno:  # comment in the same line as test/kw name
+                    header_line.append(child)
+                else:
+                    comments.append(child)
             elif getattr(child, "type", "invalid") in setting_types:
                 after_seen = after_seen or child.type in after
-                settings[child.type] = child
+                settings[child.type] = (comments, child)
+                comments = []
             elif after_seen:
+                trailing_after.extend(comments)
+                comments = []
                 trailing_after.append(child)
             else:
+                not_settings.extend(comments)
+                comments = []
                 not_settings.append(child)
+        trailing_after.extend(comments)
         # comments after last data statement are considered as comment outside body
         trailing_non_data = []
         while trailing_after and isinstance(trailing_after[-1], (EmptyLine, Comment)):
             trailing_non_data.insert(0, trailing_after.pop())
         not_settings += trailing_after
         node.body = (
-            comment
+            header_line
             + self.add_in_order(before, settings)
             + not_settings
             + self.add_in_order(after, settings)
@@ -183,4 +195,11 @@ def order_settings(self, node, setting_types, before, after):
 
     @staticmethod
     def add_in_order(order, settings_in_node):
-        return [settings_in_node[token_type] for token_type in order if token_type in settings_in_node]
+        nodes = []
+        for token_type in order:
+            if token_type not in settings_in_node:
+                continue
+            comments, node = settings_in_node[token_type]
+            nodes.extend(comments)
+            nodes.append(node)
+        return nodes

tests/atest/transformers/OrderSettings/expected/stick_comments.robot

@@ -0,0 +1,10 @@
+*** Keywords ***
+Comment between settings
+    # comment 3
+    [Documentation]    doc
+    # comment 1
+    [Arguments]    ${arg}
+    # comment 2
+    Step
+    # comment 4
+    [Return]    ${arg}

tests/atest/transformers/OrderSettings/source/stick_comments.robot

@@ -0,0 +1,10 @@
+*** Keywords ***
+Comment between settings
+    # comment 1
+    [Arguments]    ${arg}
+    # comment 2
+    Step
+    # comment 4
+    [Return]    ${arg}
+    # comment 3
+    [Documentation]    doc

tests/atest/transformers/OrderSettings/test_order_settings.py

@@ -85,3 +85,6 @@ def test_custom_order_setting_twice_in_after_before(self):
 
     def test_translated(self):
         self.compare(source="translated.robot", target_version=">=6")
+
+    def test_stick_comments_with_settings(self):
+        self.compare(source="stick_comments.robot")