PY-54512 Handle line breaks when extracting a method

GitOrigin-RevId: 95f6c067e613172913f3bbdc60eafb8623e32457
2026-09-27 10:03:11 +07:00 · 2025-09-18 15:42:03 +00:00
parent 6e0d9a581e
commit 51ad67e9f7
20 changed files with 209 additions and 7 deletions
@@ -16,6 +16,7 @@ import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.util.text.Strings;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.source.codeStyle.CodeEditUtil;
+import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.refactoring.RefactoringBundle;
 import com.intellij.refactoring.extractMethod.*;
@@ -30,10 +31,7 @@ import com.intellij.util.Function;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.ui.UIUtil;
-import com.jetbrains.python.PyNames;
-import com.jetbrains.python.PyPsiBundle;
-import com.jetbrains.python.PythonFileType;
-import com.jetbrains.python.PythonLanguage;
+import com.jetbrains.python.*;
 import com.jetbrains.python.codeInsight.codeFragment.PyCodeFragment;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
@@ -41,6 +39,7 @@ import com.jetbrains.python.codeInsight.dataflow.scope.Scope;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.codeInsight.intentions.PyTypeHintGenerationUtil;
 import com.jetbrains.python.documentation.PythonDocumentationProvider;
+import com.jetbrains.python.lexer.PythonLexer;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyFunctionBuilder;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
@@ -576,16 +575,78 @@ public final class PyExtractMethodUtil {
      builder.makeAsync();
    }
    final String text;
-    if (expression instanceof PyYieldExpression) {
-      text = String.format("(%s)", expression.getText());
+    String expressionText = expression.getText();
+    if (needToWrapInParenthesis(expression, expressionText)) {
+      text = String.format("(%s)", expressionText);
    }
    else {
-      text = expression.getText();
+      text = expressionText;
    }
    builder.statement("return " + text);
    return builder.buildFunction();
  }

+  private static boolean needToWrapInParenthesis(@NotNull PsiElement expression, @NotNull String expressionText) {
+    if (expression instanceof PyYieldExpression) {
+      return true;
+    }
+    if (expression instanceof PyGeneratorExpression) {
+      final PsiElement firstChild = expression.getFirstChild();
+      return firstChild != null && firstChild.getNode().getElementType() != PyTokenTypes.LPAR;
+    }
+    return hasLineBreakOutsideBraces(expressionText);
+  }
+
+  private static boolean hasLineBreakOutsideBraces(@NotNull String text) {
+    int parensCount = 0;
+    int bracketsCount = 0;
+    int bracesCount = 0;
+
+    PythonLexer lexer = new PythonLexer();
+    lexer.start(text);
+
+    do {
+      final IElementType tokenType = lexer.getTokenType();
+      if (tokenType == null) {
+        break;
+      }
+      if (PyTokenTypes.BACKSLASH.equals(tokenType)) {
+        lexer.advance();
+        if (PyTokenTypes.LINE_BREAK.equals(lexer.getTokenType())) {
+          lexer.advance();
+        }
+        continue;
+      }
+      if (PyTokenTypes.LPAR.equals(tokenType)) {
+        parensCount++;
+      }
+      else if (PyTokenTypes.RPAR.equals(tokenType)) {
+        parensCount--;
+      }
+      else if (PyTokenTypes.LBRACKET.equals(tokenType)) {
+        bracketsCount++;
+      }
+      else if (PyTokenTypes.RBRACKET.equals(tokenType)) {
+        bracketsCount--;
+      }
+      else if (PyTokenTypes.LBRACE.equals(tokenType)) {
+        bracesCount++;
+      }
+      else if (PyTokenTypes.RBRACE.equals(tokenType)) {
+        bracesCount--;
+      }
+      else if (PyTokenTypes.LINE_BREAK.equals(tokenType)) {
+        if (parensCount <= 0 && bracketsCount <= 0 && bracesCount <= 0) {
+          return true;
+        }
+      }
+      lexer.advance();
+    }
+    while (true);
+
+    return false;
+  }
+
  private static @NotNull PyFunction generateMethodFromElements(final @NotNull PyExtractMethodSettings methodSettings,
                                                                final @NotNull List<PsiElement> elementsRange,
                                                                @Nullable PyUtil.MethodFlags flags,
@@ -0,0 +1,8 @@
+def foo(num):
+    if bar(num):
+        return 1
+
+
+def bar(num_new):
+    return num_new >= 0 \
+        and num_new <= 9
@@ -0,0 +1,11 @@
+from typing import Any
+
+
+def foo(num):
+    if bar(num):
+        return 1
+
+
+def bar(num_new) -> Any:
+    return num_new >= 0 \
+        and num_new <= 9
@@ -0,0 +1,4 @@
+def foo(num):
+    if <selection>num >= 0 \
+       and num <= 9</selection>:
+        return 1
@@ -0,0 +1,5 @@
+def bar():
+    return (i + 1 for i in range(1, 10))
+
+
+set(bar())
@@ -0,0 +1,8 @@
+from typing import Any, Generator
+
+
+def bar() -> Generator[int, Any, None]:
+    return (i + 1 for i in range(1, 10))
+
+
+set(bar())
@@ -0,0 +1 @@
+set(<selection>i + 1 for i in range(1, 10)</selection>)
@@ -0,0 +1,6 @@
+def bar():
+    return (i + 1 for i in range(1,
+                                 10))
+
+
+_ = bar()
@@ -0,0 +1,9 @@
+from typing import Any, Generator
+
+
+def bar() -> Generator[int, Any, None]:
+    return (i + 1 for i in range(1,
+                                 10))
+
+
+_ = bar()
@@ -0,0 +1,2 @@
+_ = <selection>(i + 1 for i in range(1,
+                          10))</selection>
@@ -0,0 +1,6 @@
+def bar():
+    return [1] + [2 +
+                  3]
+
+
+_ = bar()
@@ -0,0 +1,6 @@
+def bar() -> list[int]:
+    return [1] + [2 +
+                  3]
+
+
+_ = bar()
@@ -0,0 +1,2 @@
+_ = <selection>[1] + [2 +
+           3]</selection>
@@ -0,0 +1,10 @@
+def foo(num):
+    if (
+            bar(num)
+    ):
+        return 1
+
+
+def bar(num_new):
+    return (num_new >= 0
+            and num_new <= 9)
@@ -0,0 +1,13 @@
+from typing import Any
+
+
+def foo(num):
+    if (
+            bar(num)
+    ):
+        return 1
+
+
+def bar(num_new) -> Any:
+    return (num_new >= 0
+            and num_new <= 9)
@@ -0,0 +1,6 @@
+def foo(num):
+    if (
+            <selection>num >= 0
+            and num <= 9</selection>
+    ):
+        return 1
@@ -0,0 +1,6 @@
+def bar():
+    return ([1] +
+            [2, 3])
+
+
+print(bar())
@@ -0,0 +1,6 @@
+def bar() -> list[int]:
+    return ([1] +
+            [2, 3])
+
+
+print(bar())
@@ -0,0 +1,2 @@
+print(<selection>[1] +
+      [2, 3]</selection>)
@@ -396,4 +396,34 @@ public class PyExtractMethodTest extends LightMarkedTestCase {
  public void testInnerFunctionWithSameNameAsOuterMethod() {
    doTest("foo");
  }
+
+  // PY-54512
+  public void testLineBreakOutsideBraces() {
+    doTest("bar");
+  }
+
+  // PY-54512
+  public void testLineBreakOutsideBraces2() {
+    doTest("bar");
+  }
+
+  // PY-54512
+  public void testLineBreakInsideBraces() {
+    doTest("bar");
+  }
+
+  // PY-54512
+  public void testExplicitLineJoining() {
+    doTest("bar");
+  }
+
+  // PY-54512
+  public void testGeneratorNoParens() {
+    doTest("bar");
+  }
+
+  // PY-54512
+  public void testGeneratorParenthesized() {
+    doTest("bar");
+  }
 }
				`@@ -0,0 +1 @@`
				`set(<selection>i + 1 for i in range(1, 10)</selection>)`