diff --git a/plugins/evaluation-plugin/languages/python/src/com/intellij/cce/visitor/PythonCodeGenerationVisitor.kt b/plugins/evaluation-plugin/languages/python/src/com/intellij/cce/visitor/PythonCodeGenerationVisitor.kt
index de3e680152dc..47d86b86958a 100644
--- a/plugins/evaluation-plugin/languages/python/src/com/intellij/cce/visitor/PythonCodeGenerationVisitor.kt
+++ b/plugins/evaluation-plugin/languages/python/src/com/intellij/cce/visitor/PythonCodeGenerationVisitor.kt
@@ -1,12 +1,10 @@
 package com.intellij.cce.visitor
 
 import com.intellij.cce.core.*
-import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
 import com.intellij.psi.util.startOffset
 import com.jetbrains.python.psi.PyFunction
 import com.jetbrains.python.psi.PyRecursiveElementVisitor
-import com.jetbrains.python.psi.PyStatementList
 
 class PythonCodeGenerationVisitor : CodeGenerationVisitorBase(Language.PYTHON) {
   override fun createPsiVisitor(codeFragment: CodeFragment): PsiElementVisitor {
@@ -19,16 +17,11 @@ private class PythonCodeGenerationPsiVisitor(private val codeFragment: CodeFragm
     codeFragment.addChild(
       CodeToken(node.text, node.startOffset, SimpleTokenProperties.create(TypeProperty.METHOD, SymbolLocation.PROJECT) {})
     )
+
+    val body = node.statementList
+    codeFragment.addChild(
+      CodeToken(body.text, body.startOffset, SimpleTokenProperties.create(TypeProperty.METHOD_BODY, SymbolLocation.PROJECT) {})
+    )
     super.visitPyFunction(node)
   }
-
-  override fun visitElement(element: PsiElement) {
-    if (element is PyStatementList) {
-      codeFragment.addChild(
-        CodeToken(element.text, element.startOffset, SimpleTokenProperties.create(TypeProperty.METHOD_BODY, SymbolLocation.PROJECT) {})
-      )
-    }
-
-    super.visitElement(element)
-  }
 }
\ No newline at end of file