From f1fb2a2a2deeee3878cdf9b8e364d53e6565470e Mon Sep 17 00:00:00 2001
From: Pawel Lampe <pawel.lampe@gmail.com>
Date: Sat, 14 Sep 2024 23:18:52 +0200
Subject: [PATCH] Add initial 'complex_multistatement_lambdas' tests & align
 lambda formatter implementation

---
 gdtoolkit/formatter/expression.py             | 13 +++-
 .../complex_multistatement_lambdas.in.gd      | 55 ++++++++++++++
 .../complex_multistatement_lambdas.out.gd     | 75 +++++++++++++++++++
 3 files changed, 141 insertions(+), 2 deletions(-)
 create mode 100644 tests/formatter/input-output-pairs/complex_multistatement_lambdas.in.gd
 create mode 100644 tests/formatter/input-output-pairs/complex_multistatement_lambdas.out.gd

diff --git a/gdtoolkit/formatter/expression.py b/gdtoolkit/formatter/expression.py
index abc502e0..4de7d7b2 100644
--- a/gdtoolkit/formatter/expression.py
+++ b/gdtoolkit/formatter/expression.py
@@ -772,7 +772,6 @@ def _format_lambda_to_multiple_lines(
     expression_context: ExpressionContext,
     context: Context,
 ) -> FormattedLines:
-    assert expression_context.suffix_string == ""
     expression_context_for_header = ExpressionContext(
         expression_context.prefix_string, expression_context.prefix_line, "", -1
     )
@@ -788,8 +787,18 @@ def _format_lambda_to_multiple_lines(
         function_statement_module.format_func_statement,
         child_context,
     )
+    last_block_line_number, last_block_line_content = block_lines[-1]
 
-    return header_lines + block_lines
+    return (
+        header_lines
+        + block_lines[:-1]
+        + [
+            (
+                last_block_line_number,
+                f"{last_block_line_content}{expression_context.suffix_string}",
+            )
+        ]
+    )
 
 
 def _format_lambda_header_to_multiple_lines(
diff --git a/tests/formatter/input-output-pairs/complex_multistatement_lambdas.in.gd b/tests/formatter/input-output-pairs/complex_multistatement_lambdas.in.gd
new file mode 100644
index 00000000..40018c60
--- /dev/null
+++ b/tests/formatter/input-output-pairs/complex_multistatement_lambdas.in.gd
@@ -0,0 +1,55 @@
+func foo(x):
+	pass
+func bar(x,y):
+	pass
+func baz():
+	foo(func():
+		var x = 1
+		return x)
+
+	foo(func():
+		var x = 1
+		if x > 1:
+			print(x))
+
+	foo(func():
+		var x = 1
+		if x > 1:
+			print(x)
+	)
+
+	bar(func():
+		var x = 1
+		return x, func():
+			var y = 1
+			return y)
+
+	bar(func():
+		var x = 1
+		return x,
+		func():
+			var y = 1
+			return y)
+
+	bar(func():
+		var x = 1
+		return x,
+
+		func():
+			var y = 1
+			return y)
+
+	bar(func():
+		var x = 1
+		if x > 0:
+			print(x), func():
+			var y = 1
+			return y)
+
+	bar(func():
+		var x = 1
+		if x > 0:
+			print(x),
+		func():
+			var y = 1
+			return y)
diff --git a/tests/formatter/input-output-pairs/complex_multistatement_lambdas.out.gd b/tests/formatter/input-output-pairs/complex_multistatement_lambdas.out.gd
new file mode 100644
index 00000000..6201bb28
--- /dev/null
+++ b/tests/formatter/input-output-pairs/complex_multistatement_lambdas.out.gd
@@ -0,0 +1,75 @@
+func foo(x):
+	pass
+
+
+func bar(x, y):
+	pass
+
+
+func baz():
+	foo(
+		func():
+			var x = 1
+			return x
+	)
+
+	foo(
+		func():
+			var x = 1
+			if x > 1:
+				print(x)
+	)
+
+	foo(
+		func():
+			var x = 1
+			if x > 1:
+				print(x)
+	)
+
+	bar(
+		func():
+			var x = 1
+			return x,
+		func():
+			var y = 1
+			return y
+	)
+
+	bar(
+		func():
+			var x = 1
+			return x,
+		func():
+			var y = 1
+			return y
+	)
+
+	bar(
+		func():
+			var x = 1
+			return x,
+		func():
+			var y = 1
+			return y
+	)
+
+	bar(
+		func():
+			var x = 1
+			if x > 0:
+				print(x),
+		func():
+			var y = 1
+			return y
+	)
+
+	bar(
+		func():
+			var x = 1
+			if x > 0:
+				print(x),
+		func():
+			var y = 1
+			return y
+	)