diff --git a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt index 4c59047b0387..ae5feaa41cbb 100644 --- a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt +++ b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt @@ -132,7 +132,6 @@ class PyInlineFunctionProcessor(project: Project, val functionScope = ControlFlowCache.getScope(myFunction) PyClassRefactoringUtil.rememberNamedReferences(myFunction) - val hasDocstring = myFunction.docStringExpression != null references.forEach { usage -> val reference = usage.element as PyReferenceExpression @@ -146,7 +145,11 @@ class PyInlineFunctionProcessor(project: Project, val containingStatement = PsiTreeUtil.getParentOfType(callSite, PyStatement::class.java) ?: error("Unable to find statement for ${reference.name}") val scopeAnchor = if (containingStatement is PyFunction) containingStatement else reference - val replacementFunction = myFunction.statementList.copy() as PyStatementList + val functionCopy = myFunction.copy() as PyFunction + functionCopy.typeComment?.delete() + PsiTreeUtil.getParentOfType(functionCopy.docStringExpression, PyStatement::class.java)?.delete() + + val replacement = functionCopy.statementList val namesInOuterScope = PyRefactoringUtil.collectUsedNames(refScopeOwner) val builtinCache = PyBuiltinCache.getInstance(reference) @@ -191,7 +194,7 @@ class PyInlineFunctionProcessor(project: Project, }) - replacementFunction.accept(object : PyRecursiveElementVisitor() { + replacement.accept(object : PyRecursiveElementVisitor() { override fun visitPyReferenceExpression(node: PyReferenceExpression) { if (!node.isQualified) { val parentLambda = PsiTreeUtil.getParentOfType(node, PyLambdaExpression::class.java) @@ -253,17 +256,16 @@ class PyInlineFunctionProcessor(project: Project, callSite.replace(newReturn.assignedValue!!) } - CodeStyleManager.getInstance(myProject).reformat(replacementFunction, true) + CodeStyleManager.getInstance(myProject).reformat(replacement, true) val insertElement = { elem: PsiElement -> containingStatement.parent.addBefore(elem, containingStatement) } declarations.forEach { insertElement(it) } - if (replacementFunction.firstChild != null) { - val directChildren = SyntaxTraverser.psiApi().children(replacementFunction).filter { it !is PsiWhiteSpace }.toList() - val statementsAndComments = if (hasDocstring) directChildren.drop(1) else directChildren - val statements = statementsAndComments.filterIsInstance() + if (replacement.firstChild != null) { + val children = SyntaxTraverser.psiApi().children(replacement).filter { it !is PsiWhiteSpace }.toList() + val statements = children.filterIsInstance() if (statements.size > 1 || statements.firstOrNull() !is PyPassStatement) { - statementsAndComments.asSequence() + children.asSequence() .map { insertElement(it) } .filterIsInstance() .forEach { PyClassRefactoringUtil.restoreNamedReferences(it) } diff --git a/python/testData/refactoring/inlineFunction/keepFirstComment/main.after.py b/python/testData/refactoring/inlineFunction/keepingComments/main.after.py similarity index 50% rename from python/testData/refactoring/inlineFunction/keepFirstComment/main.after.py rename to python/testData/refactoring/inlineFunction/keepingComments/main.after.py index 5b6e8cc546a2..fd61d887343b 100644 --- a/python/testData/refactoring/inlineFunction/keepFirstComment/main.after.py +++ b/python/testData/refactoring/inlineFunction/keepingComments/main.after.py @@ -1,9 +1,15 @@ def foo(): # comment on the first line + print(42) # inline comment + # comment 1 print(42) + # comment 2 return 42 # comment on the first line +print(42) # inline comment +# comment 1 print(42) +# comment 2 res = 42 \ No newline at end of file diff --git a/python/testData/refactoring/inlineFunction/keepingComments/main.py b/python/testData/refactoring/inlineFunction/keepingComments/main.py new file mode 100644 index 000000000000..87953c661103 --- /dev/null +++ b/python/testData/refactoring/inlineFunction/keepingComments/main.py @@ -0,0 +1,10 @@ +def foo(): + # comment on the first line + print(42) # inline comment + # comment 1 + print(42) + # comment 2 + return 42 + + +res = foo() \ No newline at end of file diff --git a/python/testData/refactoring/inlineFunction/removingDocstringAndTypeComment/main.after.py b/python/testData/refactoring/inlineFunction/removingDocstringAndTypeComment/main.after.py new file mode 100644 index 000000000000..0f5a57fa64a0 --- /dev/null +++ b/python/testData/refactoring/inlineFunction/removingDocstringAndTypeComment/main.after.py @@ -0,0 +1,12 @@ +def foo(): + # type: () -> int + """ + Docstring for the function + :return: int + """ + print(42) + return 42 + + +print(42) +res = 42 \ No newline at end of file diff --git a/python/testData/refactoring/inlineFunction/removingDocstringAndTypeComment/main.py b/python/testData/refactoring/inlineFunction/removingDocstringAndTypeComment/main.py new file mode 100644 index 000000000000..51fb0da9601a --- /dev/null +++ b/python/testData/refactoring/inlineFunction/removingDocstringAndTypeComment/main.py @@ -0,0 +1,11 @@ +def foo(): + # type: () -> int + """ + Docstring for the function + :return: int + """ + print(42) + return 42 + + +res = foo() \ No newline at end of file diff --git a/python/testData/refactoring/inlineFunction/removingDocstringOfEmptyFunction/main.after.py b/python/testData/refactoring/inlineFunction/removingDocstringOfEmptyFunction/main.after.py new file mode 100644 index 000000000000..1d06203ac374 --- /dev/null +++ b/python/testData/refactoring/inlineFunction/removingDocstringOfEmptyFunction/main.after.py @@ -0,0 +1,8 @@ +def foo(): + """ + This is a docstring. + :return: None + """ + pass + + diff --git a/python/testData/refactoring/inlineFunction/removingDocstringOfEmptyFunction/main.py b/python/testData/refactoring/inlineFunction/removingDocstringOfEmptyFunction/main.py new file mode 100644 index 000000000000..914cdf50a13a --- /dev/null +++ b/python/testData/refactoring/inlineFunction/removingDocstringOfEmptyFunction/main.py @@ -0,0 +1,9 @@ +def foo(): + """ + This is a docstring. + :return: None + """ + pass + + +foo() \ No newline at end of file diff --git a/python/testData/refactoring/inlineFunction/removingTypeComment/main.after.py b/python/testData/refactoring/inlineFunction/removingTypeComment/main.after.py new file mode 100644 index 000000000000..706f3dcb31de --- /dev/null +++ b/python/testData/refactoring/inlineFunction/removingTypeComment/main.after.py @@ -0,0 +1,8 @@ +def foo(): + # type: () -> int + print(42) + return 42 + + +print(42) +res = 42 \ No newline at end of file diff --git a/python/testData/refactoring/inlineFunction/keepFirstComment/main.py b/python/testData/refactoring/inlineFunction/removingTypeComment/main.py similarity index 64% rename from python/testData/refactoring/inlineFunction/keepFirstComment/main.py rename to python/testData/refactoring/inlineFunction/removingTypeComment/main.py index aaf235f31d31..a887fcd141ab 100644 --- a/python/testData/refactoring/inlineFunction/keepFirstComment/main.py +++ b/python/testData/refactoring/inlineFunction/removingTypeComment/main.py @@ -1,5 +1,5 @@ def foo(): - # comment on the first line + # type: () -> int print(42) return 42 diff --git a/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt b/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt index 6c4f746650ac..561e0f912fbc 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt +++ b/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt @@ -79,7 +79,10 @@ class PyInlineFunctionTest : PyTestCase() { fun testFunctionWithLambda() = doTest() fun testRefInDunderAll() = doTest(inlineThis = false, remove = true) fun testRemovingDocstring() = doTest() - fun testKeepFirstComment() = doTest() + fun testRemovingTypeComment() = doTest() + fun testRemovingDocstringAndTypeComment() = doTest() + fun testRemovingDocstringOfEmptyFunction() = doTest() + fun testKeepingComments() = doTest() //fun testInlineImportedAs() = doTest(inlineThis = false) fun testRemoveFunctionWithStub() { doTest(inlineThis = false, remove = true)