IDEA-CR-49176: function inline: remove docstring, keep the first comment, unless it's a type comment(PY-36714, PY-36491)

GitOrigin-RevId: d7de37f5a7b3b62cc6015e2ff3dbdad5cafe3641

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<PyStatement>()
+      if (replacement.firstChild != null) {
+        val children = SyntaxTraverser.psiApi().children(replacement).filter { it !is PsiWhiteSpace }.toList()
+        val statements = children.filterIsInstance<PyStatement>()
         if (statements.size > 1 || statements.firstOrNull() !is PyPassStatement) {
-          statementsAndComments.asSequence()
+          children.asSequence()
             .map { insertElement(it) }
             .filterIsInstance<PyStatement>()
             .forEach { PyClassRefactoringUtil.restoreNamedReferences(it) }

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

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 = fo<caret>o()

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

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 = fo<caret>o()

python/testData/refactoring/inlineFunction/removingDocstringOfEmptyFunction/main.after.py

@@ -0,0 +1,8 @@
+def foo():
+    """
+    This is a docstring.
+    :return: None
+    """
+    pass
+
+

python/testData/refactoring/inlineFunction/removingDocstringOfEmptyFunction/main.py

@@ -0,0 +1,9 @@
+def foo():
+    """
+    This is a docstring.
+    :return: None
+    """
+    pass
+
+
+fo<caret>o()

python/testData/refactoring/inlineFunction/removingTypeComment/main.after.py

@@ -0,0 +1,8 @@
+def foo():
+    # type: () -> int
+    print(42)
+    return 42
+
+
+print(42)
+res = 42

python/testData/refactoring/inlineFunction/removingTypeComment/main.py

@@ -1,5 +1,5 @@
 def foo():
-    # comment on the first line
+    # type: () -> int
     print(42)
     return 42
 

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)