IDEA-CR-49176: inline function doesn't copy docstring, but keeps the first-line comment (PY-36469, PY-36491)

GitOrigin-RevId: db1fbb5178892db955678dfcec684a5af7f987c6

python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt

@@ -121,6 +121,7 @@ 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
@@ -231,11 +232,19 @@ class PyInlineFunctionProcessor(project: Project,
 
       CodeStyleManager.getInstance(myProject).reformat(replacementFunction, true)
 
-      declarations.forEach { containingStatement.parent.addBefore(it, containingStatement) }
+      val insertElement = { elem: PsiElement -> containingStatement.parent.addBefore(elem, containingStatement) }
+
+      declarations.forEach { insertElement(it) }
       if (replacementFunction.firstChild != null) {
-        val statements = replacementFunction.statements
+        if (replacementFunction.firstChild is PsiComment) {
+          SyntaxTraverser.psiApi().children(replacementFunction)
+            .takeWhile { it is PsiComment || it is PsiWhiteSpace }
+            .filterIsInstance<PsiComment>()
+            .forEach { insertElement(it) }
+        }
+        val statements = if (hasDocstring) replacementFunction.statements.drop(1) else replacementFunction.statements.toList()
         statements.asSequence()
-          .map { containingStatement.parent.addBefore(it, containingStatement) }
+          .map { insertElement(it) }
           .forEach { PyClassRefactoringUtil.restoreNamedReferences(it) }
       }
 

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

@@ -0,0 +1,9 @@
+def foo():
+    # comment on the first line
+    print(42)
+    return 42
+
+
+# comment on the first line
+print(42)
+res = 42

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

@@ -0,0 +1,7 @@
+def foo():
+    # comment on the first line
+    print(42)
+    return 42
+
+
+res = fo<caret>o()

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

@@ -0,0 +1,11 @@
+def foo():
+    """
+    Docstring for the function
+    :return: int
+    """
+    print(42)
+    return 42
+
+
+print(42)
+res = 42

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

@@ -0,0 +1,10 @@
+def foo():
+    """
+    Docstring for the function
+    :return: int
+    """
+    print(42)
+    return 42
+
+
+res = fo<caret>o()

python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt

@@ -73,6 +73,8 @@ class PyInlineFunctionTest : PyTestCase() {
   fun testCallFromStaticMethod() = doTest()
   fun testCallFromClassMethod() = doTest()
   fun testComplexQualifier() = doTest()
+  fun testRemovingDocstring() = doTest()
+  fun testKeepFirstComment() = doTest()
   //fun testInlineImportedAs() = doTest(inlineThis = false)
   fun testRemoveFunctionWithStub() {
     doTest(inlineThis = false, remove = true)