add extra newline to execute when returning dictionary

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

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