Hang at equals token in table field if comments before value (#982)

* Add test case

* Hang after equals in table if comments present

* Update changelog

* Update snapshots

* Add another test case

* Clippy

Author: JohnnyMorganz

CHANGELOG.md

@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed panic when attempting to format a file outside of the current working directory when `--respect-ignores` is enabled ([#969](https://github.com/JohnnyMorganz/StyLua/pull/969))
 - Fixed unnecessary semicolons being introduced at the end of statements when incorrectly determined as ambiguous ([#963](https://github.com/JohnnyMorganz/StyLua/issues/963))
 - Fixed malformed formatting of function calls where parentheses are removed but there are comments in between the parentheses and the expression. Now, we will keep the parentheses in these cases, except for trailing comments ([#964](https://github.com/JohnnyMorganz/StyLua/issues/964))
+- Fixed malformed formatting of table field expression when there are comments in between the equals and the value ([#942](https://github.com/JohnnyMorganz/StyLua/issues/942))
 
 ## [2.0.2] - 2024-12-07
 

src/formatters/assignment.rs

@@ -162,6 +162,48 @@ fn prevent_equals_hanging(expression: &Expression) -> bool {
     }
 }
 
+pub fn hang_at_equals_due_to_comments(
+    ctx: &Context,
+    equal_token: &TokenReference,
+    expression: &Expression,
+    shape: Shape,
+) -> (TokenReference, Expression) {
+    // We will hang at the equals token, and then format the expression as necessary
+    let equal_token = hang_equal_token(ctx, equal_token, shape, false);
+
+    let shape = shape.reset().increment_additional_indent();
+
+    // As we know that there is only a single element in the list, we can extract it to work with it
+    // Format the expression given - if it contains comments, make sure to hang the expression
+    // Ignore the leading comments though (as they are solved by hanging at the equals), and the
+    // trailing comments, as they don't affect anything
+    let expression = if strip_trivia(expression).has_inline_comments() {
+        hang_expression(ctx, expression, shape, None)
+    } else {
+        format_expression(ctx, expression, shape)
+    };
+
+    // We need to take all the leading trivia from the expr_list
+    let (expression, leading_comments) = trivia_util::take_leading_comments(&expression);
+
+    // Indent each comment and trail them with a newline
+    let leading_comments = leading_comments
+        .iter()
+        .flat_map(|x| {
+            vec![
+                create_indent_trivia(ctx, shape),
+                x.to_owned(),
+                create_newline_trivia(ctx),
+            ]
+        })
+        .chain(std::iter::once(create_indent_trivia(ctx, shape)))
+        .collect();
+
+    let expression = expression.update_leading_trivia(FormatTriviaType::Replace(leading_comments));
+
+    (equal_token, expression)
+}
+
 /// Attempts different formatting tactics on an expression list being assigned (`= foo, bar`), to find the best
 /// formatting output.
 fn attempt_assignment_tactics(
@@ -240,39 +282,8 @@ fn attempt_assignment_tactics(
         if trivia_util::token_contains_comments(&equal_token)
             || expression.has_leading_comments(CommentSearch::Single)
         {
-            // We will hang at the equals token, and then format the expression as necessary
-            let equal_token = hang_equal_token(ctx, &equal_token, shape, false);
-
-            let shape = shape.reset().increment_additional_indent();
-
-            // As we know that there is only a single element in the list, we can extract it to work with it
-            // Format the expression given - if it contains comments, make sure to hang the expression
-            // Ignore the leading comments though (as they are solved by hanging at the equals), and the
-            // trailing comments, as they don't affect anything
-            let expression = if strip_trivia(expression).has_inline_comments() {
-                hang_expression(ctx, expression, shape, None)
-            } else {
-                format_expression(ctx, expression, shape)
-            };
-
-            // We need to take all the leading trivia from the expr_list
-            let (expression, leading_comments) = trivia_util::take_leading_comments(&expression);
-
-            // Indent each comment and trail them with a newline
-            let leading_comments = leading_comments
-                .iter()
-                .flat_map(|x| {
-                    vec![
-                        create_indent_trivia(ctx, shape),
-                        x.to_owned(),
-                        create_newline_trivia(ctx),
-                    ]
-                })
-                .chain(std::iter::once(create_indent_trivia(ctx, shape)))
-                .collect();
-
-            let expression =
-                expression.update_leading_trivia(FormatTriviaType::Replace(leading_comments));
+            let (equal_token, expression) =
+                hang_at_equals_due_to_comments(ctx, &equal_token, expression, shape);
 
             // Rebuild expression back into a list
             let expr_list = std::iter::once(Pair::new(expression, None)).collect();

src/formatters/table.rs

@@ -23,6 +23,8 @@ use full_moon::{
     tokenizer::{Token, TokenReference, TokenType},
 };
 
+use super::{assignment::hang_at_equals_due_to_comments, trivia_util::GetLeadingTrivia};
+
 /// Used to provide information about the table
 #[derive(Debug, Clone, Copy)]
 pub enum TableType {
@@ -68,6 +70,25 @@ fn format_field_expression_value(
     }
 }
 
+fn hang_field_value_at_equals_due_to_comments(
+    ctx: &Context,
+    equal: &TokenReference,
+    expression: &Expression,
+    shape: Shape,
+) -> (TokenReference, Expression) {
+    // Remove singleline comments from the output expression as it will be moved after the comma
+    // Retain multiline comments in place
+    let multiline_comments = expression.trailing_comments_search(CommentSearch::Multiline);
+    let trailing_trivia = FormatTriviaType::Replace(multiline_comments);
+
+    let (equal_token, expression) = hang_at_equals_due_to_comments(ctx, equal, expression, shape);
+
+    (
+        equal_token,
+        expression.update_trailing_trivia(trailing_trivia),
+    )
+}
+
 /// Handles the formatting of the comments around a key and the equals sign in a field of a table.
 /// Takes in the key as a node (so that we can handle both expression key brackets and name keys)
 /// as well as the equals sign, then outputs the new leading trivia of the key + new equals token. The trailing trivia of the key should be emptied.
@@ -177,8 +198,14 @@ fn format_field(
                 .update_trailing_trivia(FormatTriviaType::Replace(vec![]))
                 .update_leading_trivia(leading_trivia);
 
-            let shape = shape.take_last_line(&key) + (2 + 3 + if space_brackets { 2 } else { 0 }); // 2 = brackets, 3 = " = ", 2 = spaces around brackets if necessary
-            let value = format_field_expression_value(ctx, value, shape);
+            let (equal, value) = if value.has_leading_comments(CommentSearch::Single) {
+                hang_field_value_at_equals_due_to_comments(ctx, &equal, value, shape)
+            } else {
+                let shape =
+                    shape.take_last_line(&key) + (2 + 3 + if space_brackets { 2 } else { 0 }); // 2 = brackets, 3 = " = ", 2 = spaces around brackets if necessary
+                let value = format_field_expression_value(ctx, value, shape);
+                (equal, value)
+            };
 
             Field::ExpressionKey {
                 brackets,
@@ -202,8 +229,13 @@ fn format_field(
                 .update_trailing_trivia(FormatTriviaType::Replace(vec![]))
                 .update_leading_trivia(leading_trivia);
 
-            let shape = shape + (strip_trivia(&key).to_string().len() + 3); // 3 = " = "
-            let value = format_field_expression_value(ctx, value, shape);
+            let (equal, value) = if value.has_leading_comments(CommentSearch::Single) {
+                hang_field_value_at_equals_due_to_comments(ctx, &equal, value, shape)
+            } else {
+                let shape = shape + (strip_trivia(&key).to_string().len() + 3); // 3 = " = "
+                let value = format_field_expression_value(ctx, value, shape);
+                (equal, value)
+            };
 
             Field::NameKey { key, equal, value }
         }

tests/inputs/table-field-comments-2.lua

@@ -0,0 +1,9 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/942
+local t = {
+	plus_one =
+    ---@param n number
+	---@return number
+    function(n)
+		return n + 1
+	end  ,
+}

tests/inputs/table-field-comments-3.lua

@@ -0,0 +1,15 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/942
+local mixed = {
+	--- identity
+	---@param n number
+	---@return number
+	function(n)
+		return n
+	end,
+	plus_one =
+		---@param n number
+		---@return number
+		function(n)
+			return n + 1
+		end,
+}

tests/snapshots/[email protected]

@@ -0,0 +1,15 @@
+---
+source: tests/tests.rs
+expression: "format(&contents, LuaVersion::Lua51)"
+input_file: tests/inputs/table-field-comments-2.lua
+snapshot_kind: text
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/942
+local t = {
+	plus_one =
+		---@param n number
+		---@return number
+		function(n)
+			return n + 1
+		end,
+}

tests/snapshots/[email protected]

@@ -0,0 +1,21 @@
+---
+source: tests/tests.rs
+expression: "format(&contents, LuaVersion::Lua51)"
+input_file: tests/inputs/table-field-comments-3.lua
+snapshot_kind: text
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/942
+local mixed = {
+	--- identity
+	---@param n number
+	---@return number
+	function(n)
+		return n
+	end,
+	plus_one =
+		---@param n number
+		---@return number
+		function(n)
+			return n + 1
+		end,
+}