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