Fix --flattenkeywords with VAR, GROUP and RETURN

pekkaklarck · pekkaklarck · commit 0c9cff3abcff · 2025-07-02T22:13:34.000+03:00
Fixes #5466.
diff --git a/atest/robot/output/flatten_keyword.robot b/atest/robot/output/flatten_keyword.robot
@@ -68,14 +68,14 @@ Flattened in log after execution
 
 Flatten controls in keyword
     ${tc} =    Check Test Case    ${TEST NAME}
-    Check Counts        ${tc[0]}              23
     @{expected} =    Create List
     ...    Outside IF    Inside IF    1    Nested IF
     ...    3    2    1    BANG!
     ...    FOR: 0    1    FOR: 1    1    FOR: 2    1
     ...    WHILE: 2    1    \${i} = 1    WHILE: 1    1    \${i} = 0
     ...    AssertionError    1    finally
-    FOR    ${msg}    ${exp}    IN ZIP    ${tc[0].body}    ${expected}
+    ...    Inside GROUP    \${x} = Using VAR
+    FOR    ${msg}    ${exp}    IN ZIP    ${tc[0].body}    ${expected}    mode=STRICT
         Check Log Message    ${msg}    ${exp}    level=IGNORE
     END
 
diff --git a/atest/testdata/output/flatten_keywords.robot b/atest/testdata/output/flatten_keywords.robot
@@ -88,6 +88,8 @@ Flatten controls in keyword
     FOR    ${i}    IN RANGE    3
         Log   FOR: ${i}
         Keyword 1
+        CONTINUE
+        Fail    Not run
     END
     WHILE    $i > 0
         Log   WHILE: ${i}
@@ -103,6 +105,11 @@ Flatten controls in keyword
     FINALLY
         Log    finally
     END
+    GROUP
+        Log    Inside GROUP
+    END
+    VAR    ${x}    Using VAR
+    RETURN    return value
 
 Countdown
     [Arguments]    ${count}=${3}
diff --git a/src/robot/result/resultbuilder.py b/src/robot/result/resultbuilder.py
@@ -172,27 +172,31 @@ def _flatten_keywords(self, context, flattened):
         name_match, by_name = self._get_matcher(FlattenByNameMatcher, flattened)
         type_match, by_type = self._get_matcher(FlattenByTypeMatcher, flattened)
         started = -1  # If 0 or more, we are flattening.
-        containers = {"kw", "for", "while", "iter", "if", "try"}
+        include_on_top = {"doc", "tag", "timeout", "status"}
         for event, elem in context:
             tag = elem.tag
             if event == "start":
-                if tag in containers:
-                    if started >= 0:
-                        started += 1
-                    elif by_name and tag == "kw" and name_match(
+                if started >= 0:
+                    started += 1
+                elif (
+                    by_name
+                    and tag == "kw"
+                    and name_match(
                         elem.get("name", ""),
                         elem.get("owner") or elem.get("library"),
-                    ):
-                        started = 0
-                    elif by_type and type_match(tag):
-                        started = 0
-            elif started == 0 and tag == "status":
+                        # 'library' is for RF < 7 compatibility
+                    )
+                ):
+                    started = 0
+                elif by_type and type_match(tag):
+                    started = 0
+            elif started == 1 and tag == "status":
                 elem.text = create_flatten_message(elem.text)
-            if started <= 0 or tag == "msg":
+            if started <= 0 or (started == 1 and tag in include_on_top) or tag == "msg":
                 yield event, elem
             else:
                 elem.clear()
-            if started >= 0 and event == "end" and tag in containers:
+            if started >= 0 and event == "end":
                 started -= 1
 
     def _get_matcher(self, matcher_class, flattened):
