From d5b19e7be9ce1e5d9d4b3f8c78d2f6b5df76c29a Mon Sep 17 00:00:00 2001
From: Anthony Kim <62267334+anthonykim1@users.noreply.github.com>
Date: Wed, 5 Feb 2025 13:51:17 -0800
Subject: [PATCH] add extra newline to execute when returning dictionary
 (#24784)

Resolves: https://github.com/microsoft/vscode-python/issues/22469
Need one more extra line to "execute" on behalf of user when returning
dictionary.
---
 python_files/normalizeSelection.py            |  9 ++++
 .../tests/test_normalize_selection.py         | 47 +++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/python_files/normalizeSelection.py b/python_files/normalizeSelection.py
index 3d5137fe4aeb..9d82a4dc9440 100644
--- a/python_files/normalizeSelection.py
+++ b/python_files/normalizeSelection.py
@@ -120,8 +120,12 @@ def normalize_lines(selection):
 
         # Insert a newline between each top-level statement, and append a newline to the selection.
         source = "\n".join(statements) + "\n"
+        # If selection ends with trailing dictionary or list, remove last unnecessary newline.
         if selection[-2] == "}" or selection[-2] == "]":
             source = source[:-1]
+        # If the selection contains trailing return dictionary, insert newline to trigger execute.
+        if check_end_with_return_dict(selection):
+            source = source + "\n"
     except Exception:
         # If there's a problem when parsing statements,
         # append a blank line to end the block and send it as-is.
@@ -134,6 +138,11 @@ def normalize_lines(selection):
 min_key = None
 
 
+def check_end_with_return_dict(code):
+    stripped_code = code.strip()
+    return stripped_code.endswith("}") and "return {" in stripped_code.strip()
+
+
 def check_exact_exist(top_level_nodes, start_line, end_line):
     return [
         node
diff --git a/python_files/tests/test_normalize_selection.py b/python_files/tests/test_normalize_selection.py
index e16eb118db12..779bb9720bfa 100644
--- a/python_files/tests/test_normalize_selection.py
+++ b/python_files/tests/test_normalize_selection.py
@@ -268,3 +268,50 @@ def test_list_comp(self):
         result = normalizeSelection.normalize_lines(src)
 
         assert result == expected
+
+    def test_return_dict(self):
+        importlib.reload(normalizeSelection)
+        src = textwrap.dedent(
+            """\
+            def get_dog(name, breed):
+                return {'name': name, 'breed': breed}
+            """
+        )
+
+        expected = textwrap.dedent(
+            """\
+            def get_dog(name, breed):
+                return {'name': name, 'breed': breed}
+
+            """
+        )
+
+        result = normalizeSelection.normalize_lines(src)
+
+        assert result == expected
+
+    def test_return_dict2(self):
+        importlib.reload(normalizeSelection)
+        src = textwrap.dedent(
+            """\
+            def get_dog(name, breed):
+                return {'name': name, 'breed': breed}
+
+            dog = get_dog('Ahri', 'Pomeranian')
+            print(dog)
+            """
+        )
+
+        expected = textwrap.dedent(
+            """\
+            def get_dog(name, breed):
+                return {'name': name, 'breed': breed}
+
+            dog = get_dog('Ahri', 'Pomeranian')
+            print(dog)
+            """
+        )
+
+        result = normalizeSelection.normalize_lines(src)
+
+        assert result == expected