From 51ad67e9f7ac67dacc03369de4cfbc84f7d2053a Mon Sep 17 00:00:00 2001 From: Petr Date: Wed, 17 Sep 2025 14:43:51 +0200 Subject: [PATCH] PY-54512 Handle line breaks when extracting a method GitOrigin-RevId: 95f6c067e613172913f3bbdc60eafb8623e32457 --- .../extractmethod/PyExtractMethodUtil.java | 75 +++++++++++++++++-- .../ExplicitLineJoining.after.py | 8 ++ .../ExplicitLineJoining.after.withTypes.py | 11 +++ .../ExplicitLineJoining.before.py | 4 + .../extractmethod/GeneratorNoParens.after.py | 5 ++ .../GeneratorNoParens.after.withTypes.py | 8 ++ .../extractmethod/GeneratorNoParens.before.py | 1 + .../GeneratorParenthesized.after.py | 6 ++ .../GeneratorParenthesized.after.withTypes.py | 9 +++ .../GeneratorParenthesized.before.py | 2 + .../LineBreakInsideBraces.after.py | 6 ++ .../LineBreakInsideBraces.after.withTypes.py | 6 ++ .../LineBreakInsideBraces.before.py | 2 + .../LineBreakOutsideBraces.after.py | 10 +++ .../LineBreakOutsideBraces.after.withTypes.py | 13 ++++ .../LineBreakOutsideBraces.before.py | 6 ++ .../LineBreakOutsideBraces2.after.py | 6 ++ ...LineBreakOutsideBraces2.after.withTypes.py | 6 ++ .../LineBreakOutsideBraces2.before.py | 2 + .../refactoring/PyExtractMethodTest.java | 30 ++++++++ 20 files changed, 209 insertions(+), 7 deletions(-) create mode 100644 python/testData/refactoring/extractmethod/ExplicitLineJoining.after.py create mode 100644 python/testData/refactoring/extractmethod/ExplicitLineJoining.after.withTypes.py create mode 100644 python/testData/refactoring/extractmethod/ExplicitLineJoining.before.py create mode 100644 python/testData/refactoring/extractmethod/GeneratorNoParens.after.py create mode 100644 python/testData/refactoring/extractmethod/GeneratorNoParens.after.withTypes.py create mode 100644 python/testData/refactoring/extractmethod/GeneratorNoParens.before.py create mode 100644 python/testData/refactoring/extractmethod/GeneratorParenthesized.after.py create mode 100644 python/testData/refactoring/extractmethod/GeneratorParenthesized.after.withTypes.py create mode 100644 python/testData/refactoring/extractmethod/GeneratorParenthesized.before.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakInsideBraces.after.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakInsideBraces.after.withTypes.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakInsideBraces.before.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakOutsideBraces.after.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakOutsideBraces.after.withTypes.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakOutsideBraces.before.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.after.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.after.withTypes.py create mode 100644 python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.before.py diff --git a/python/python-psi-impl/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java b/python/python-psi-impl/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java index 735f8e7df962..6335351ed01c 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java +++ b/python/python-psi-impl/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java @@ -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 elementsRange, @Nullable PyUtil.MethodFlags flags, diff --git a/python/testData/refactoring/extractmethod/ExplicitLineJoining.after.py b/python/testData/refactoring/extractmethod/ExplicitLineJoining.after.py new file mode 100644 index 000000000000..cc6e3b29abbc --- /dev/null +++ b/python/testData/refactoring/extractmethod/ExplicitLineJoining.after.py @@ -0,0 +1,8 @@ +def foo(num): + if bar(num): + return 1 + + +def bar(num_new): + return num_new >= 0 \ + and num_new <= 9 \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/ExplicitLineJoining.after.withTypes.py b/python/testData/refactoring/extractmethod/ExplicitLineJoining.after.withTypes.py new file mode 100644 index 000000000000..cd6845b97ff8 --- /dev/null +++ b/python/testData/refactoring/extractmethod/ExplicitLineJoining.after.withTypes.py @@ -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 \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/ExplicitLineJoining.before.py b/python/testData/refactoring/extractmethod/ExplicitLineJoining.before.py new file mode 100644 index 000000000000..cd64f80b3de5 --- /dev/null +++ b/python/testData/refactoring/extractmethod/ExplicitLineJoining.before.py @@ -0,0 +1,4 @@ +def foo(num): + if num >= 0 \ + and num <= 9: + return 1 \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/GeneratorNoParens.after.py b/python/testData/refactoring/extractmethod/GeneratorNoParens.after.py new file mode 100644 index 000000000000..c68c3ff10f4c --- /dev/null +++ b/python/testData/refactoring/extractmethod/GeneratorNoParens.after.py @@ -0,0 +1,5 @@ +def bar(): + return (i + 1 for i in range(1, 10)) + + +set(bar()) \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/GeneratorNoParens.after.withTypes.py b/python/testData/refactoring/extractmethod/GeneratorNoParens.after.withTypes.py new file mode 100644 index 000000000000..f48b5a159efa --- /dev/null +++ b/python/testData/refactoring/extractmethod/GeneratorNoParens.after.withTypes.py @@ -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()) \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/GeneratorNoParens.before.py b/python/testData/refactoring/extractmethod/GeneratorNoParens.before.py new file mode 100644 index 000000000000..112b2a4f0921 --- /dev/null +++ b/python/testData/refactoring/extractmethod/GeneratorNoParens.before.py @@ -0,0 +1 @@ +set(i + 1 for i in range(1, 10)) \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/GeneratorParenthesized.after.py b/python/testData/refactoring/extractmethod/GeneratorParenthesized.after.py new file mode 100644 index 000000000000..3650751b4373 --- /dev/null +++ b/python/testData/refactoring/extractmethod/GeneratorParenthesized.after.py @@ -0,0 +1,6 @@ +def bar(): + return (i + 1 for i in range(1, + 10)) + + +_ = bar() \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/GeneratorParenthesized.after.withTypes.py b/python/testData/refactoring/extractmethod/GeneratorParenthesized.after.withTypes.py new file mode 100644 index 000000000000..4cfcb25502ea --- /dev/null +++ b/python/testData/refactoring/extractmethod/GeneratorParenthesized.after.withTypes.py @@ -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() \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/GeneratorParenthesized.before.py b/python/testData/refactoring/extractmethod/GeneratorParenthesized.before.py new file mode 100644 index 000000000000..c5b03a232396 --- /dev/null +++ b/python/testData/refactoring/extractmethod/GeneratorParenthesized.before.py @@ -0,0 +1,2 @@ +_ = (i + 1 for i in range(1, + 10)) \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakInsideBraces.after.py b/python/testData/refactoring/extractmethod/LineBreakInsideBraces.after.py new file mode 100644 index 000000000000..fa65aa60eb99 --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakInsideBraces.after.py @@ -0,0 +1,6 @@ +def bar(): + return [1] + [2 + + 3] + + +_ = bar() \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakInsideBraces.after.withTypes.py b/python/testData/refactoring/extractmethod/LineBreakInsideBraces.after.withTypes.py new file mode 100644 index 000000000000..6387fa44344d --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakInsideBraces.after.withTypes.py @@ -0,0 +1,6 @@ +def bar() -> list[int]: + return [1] + [2 + + 3] + + +_ = bar() \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakInsideBraces.before.py b/python/testData/refactoring/extractmethod/LineBreakInsideBraces.before.py new file mode 100644 index 000000000000..b4db38ac9b86 --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakInsideBraces.before.py @@ -0,0 +1,2 @@ +_ = [1] + [2 + + 3] \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.after.py b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.after.py new file mode 100644 index 000000000000..659b52c82b36 --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.after.py @@ -0,0 +1,10 @@ +def foo(num): + if ( + bar(num) + ): + return 1 + + +def bar(num_new): + return (num_new >= 0 + and num_new <= 9) \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.after.withTypes.py b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.after.withTypes.py new file mode 100644 index 000000000000..ccccb86d1f77 --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.after.withTypes.py @@ -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) \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.before.py b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.before.py new file mode 100644 index 000000000000..2055952d89a4 --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces.before.py @@ -0,0 +1,6 @@ +def foo(num): + if ( + num >= 0 + and num <= 9 + ): + return 1 \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.after.py b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.after.py new file mode 100644 index 000000000000..e076cd7e9d59 --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.after.py @@ -0,0 +1,6 @@ +def bar(): + return ([1] + + [2, 3]) + + +print(bar()) \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.after.withTypes.py b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.after.withTypes.py new file mode 100644 index 000000000000..6fad40f66fb1 --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.after.withTypes.py @@ -0,0 +1,6 @@ +def bar() -> list[int]: + return ([1] + + [2, 3]) + + +print(bar()) \ No newline at end of file diff --git a/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.before.py b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.before.py new file mode 100644 index 000000000000..e058e5cd1646 --- /dev/null +++ b/python/testData/refactoring/extractmethod/LineBreakOutsideBraces2.before.py @@ -0,0 +1,2 @@ +print([1] + + [2, 3]) \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/refactoring/PyExtractMethodTest.java b/python/testSrc/com/jetbrains/python/refactoring/PyExtractMethodTest.java index 6f42f6f9a6da..234a07ed8925 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/PyExtractMethodTest.java +++ b/python/testSrc/com/jetbrains/python/refactoring/PyExtractMethodTest.java @@ -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"); + } }