Fix selection normalization w/ decorators

news/2 Fixes/15058.md

@@ -0,0 +1 @@
+Handle decorators properly when using the `Run Selection/Line in Python Terminal` command.

pythonFiles/normalizeSelection.py

@@ -46,7 +46,19 @@ def _get_statements(selection):
     # Therefore, to retrieve the end line of each block in a version-agnostic way we need to do
     # `end = next_block.lineno - 1`
     # for all blocks except the last one, which will will just run until the last line.
-    ends = [node.lineno - 1 for node in tree.body[1:]] + [len(lines)]
+    ends = []
+    for node in tree.body[1:]:
+        line_end = node.lineno - 1
+        # Special handling of decorators:
+        # In Python 3, decorators are not taken into account in the value returned by lineno,
+        # and we have to use the length of the decorator_list array to compute the actual start line.
+        # In Python 2.7, lineno takes into account decorators, so this offset check is unnecessary.
+        # Also, not all AST objects can have decorators.
+        if hasattr(node, "decorator_list") and sys.version_info.major >= 3:
+            line_end -= len(node.decorator_list)
+        ends.append(line_end)
+    ends.append(len(lines))
+
     for node, end in zip(tree.body, ends):
         # Given this selection:
         # 1: if (m > 0 and
@@ -56,6 +68,10 @@ def _get_statements(selection):
         #
         # The first block would have lineno = 1,and the second block lineno = 4
         start = node.lineno - 1
+
+        # Special handling of decorators similar to what's above.
+        if hasattr(node, "decorator_list") and sys.version_info.major >= 3:
+            start -= len(node.decorator_list)
         block = "\n".join(lines[start:end])
 
         # If the block is multiline, add an extra newline character at its end.

pythonFiles/tests/test_normalize_selection.py

@@ -28,7 +28,7 @@ def show_something():
             """\
             def show_something():
                 print("Something")
-            
+
             """
         )
         result = normalizeSelection.normalize_lines(src)
@@ -54,7 +54,7 @@ def test_withHangingIndent(self):
             result = x + y + z
             if result == 42:
                 print("The answer to life, the universe, and everything")
-            
+
             """
         )
         result = normalizeSelection.normalize_lines(src)
@@ -119,7 +119,7 @@ def test_partialSingleLine(self):
 
     def test_multiLineWithIndent(self):
         src = """\
-           
+
         if (x > 0
             and condition == True):
             print('foo')
@@ -135,7 +135,7 @@ def test_multiLineWithIndent(self):
             print('foo')
         else:
             print('bar')
-        
+
         """
         )
 
@@ -156,7 +156,7 @@ def show_something():
             def show_something():
                 # A comment
                 print("Something")
-            
+
             """
         )
         result = normalizeSelection.normalize_lines(src)
@@ -179,3 +179,39 @@ def show_something():
         expected = src + "\n\n"
         result = normalizeSelection.normalize_lines(src)
         assert result == expected
+
+    def test_decorators(self):
+        src = textwrap.dedent(
+            """\
+            def foo(func):
+
+                def wrapper():
+                    print('before')
+                    func()
+                    print('after')
+
+                return wrapper
+
+
+            @foo
+            def show_something():
+                print("Something")
+            """
+        )
+        expected = textwrap.dedent(
+            """\
+            def foo(func):
+                def wrapper():
+                    print('before')
+                    func()
+                    print('after')
+                return wrapper
+
+            @foo
+            def show_something():
+                print("Something")
+
+            """
+        )
+        result = normalizeSelection.normalize_lines(src)
+        assert result == expected