From cb1875f3c65c0d877c8aa6bb34cf9252ef03ff65 Mon Sep 17 00:00:00 2001
From: Aleksei Kniazev <aleksei.kniazev@jetbrains.com>
Date: Thu, 27 Jun 2019 18:37:57 +0300
Subject: [PATCH] IDEA-CR-49176: all comments are moved with the function body
 after inline (PY-36491)

GitOrigin-RevId: b25b35392252d9855f23f98381505c5de5c948d5
---
 .../refactoring/inline/PyInlineFunctionProcessor.kt | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
index c6fbb6321c4c..960658a1ba06 100644
--- a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
+++ b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
@@ -255,16 +255,13 @@ class PyInlineFunctionProcessor(project: Project,
 
       declarations.forEach { insertElement(it) }
       if (replacementFunction.firstChild != null) {
-        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()
+        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 (statements.size > 1 || statements.firstOrNull() !is PyPassStatement) {
-          statements.asSequence()
+          statementsAndComments.asSequence()
             .map { insertElement(it) }
+            .filterIsInstance<PyStatement>()
             .forEach { PyClassRefactoringUtil.restoreNamedReferences(it) }
         }
       }