From 0ea52e454672541d5ae0a7489bddd9623364711e Mon Sep 17 00:00:00 2001
From: Anthony Kim <anthonykim@microsoft.com>
Date: Wed, 5 Feb 2025 08:52:21 -0800
Subject: [PATCH 1/3] add extra newline to execute for return dictionary

---
 python_files/normalizeSelection.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/python_files/normalizeSelection.py b/python_files/normalizeSelection.py
index 3d5137fe4aeb..0cfd3343a4e3 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 new-line for "cleaner run".
         if selection[-2] == "}" or selection[-2] == "]":
             source = source[:-1]
+        # If the selection contains trailing return dictionary, insert newline to run automatically.
+        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

From 641993060cc84c5f4dc479397f75e56753f4c7fe Mon Sep 17 00:00:00 2001
From: Anthony Kim <anthonykim@microsoft.com>
Date: Wed, 5 Feb 2025 09:14:45 -0800
Subject: [PATCH 2/3] add tests

---
 .../tests/test_normalize_selection.py         | 47 +++++++++++++++++++
 1 file changed, 47 insertions(+)

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

From 971599d7c771ecd2045e828f0fdf4b8bbfc22b46 Mon Sep 17 00:00:00 2001
From: Anthony Kim <anthonykim@microsoft.com>
Date: Wed, 5 Feb 2025 09:28:12 -0800
Subject: [PATCH 3/3] detailed comments

---
 python_files/normalizeSelection.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python_files/normalizeSelection.py b/python_files/normalizeSelection.py
index 0cfd3343a4e3..9d82a4dc9440 100644
--- a/python_files/normalizeSelection.py
+++ b/python_files/normalizeSelection.py
@@ -120,10 +120,10 @@ 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 new-line for "cleaner run".
+        # 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 run automatically.
+        # If the selection contains trailing return dictionary, insert newline to trigger execute.
         if check_end_with_return_dict(selection):
             source = source + "\n"
     except Exception: