1 files changed, 5 insertions, 12 deletions
diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/enterProcessors/PyPlainEnterProcessor.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/enterProcessors/PyPlainEnterProcessor.java
index ac2e31252dcf..ccdcf1b93f2a 100644
--- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/enterProcessors/PyPlainEnterProcessor.java
+++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/enterProcessors/PyPlainEnterProcessor.java
@@ -20,7 +20,9 @@ import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.codeInsight.editorActions.smartEnter.SmartEnterUtil;
-import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.PyStatementList;
+import com.jetbrains.python.psi.PyStatementListContainer;
+import com.jetbrains.python.psi.PyStatementPart;
 import org.jetbrains.annotations.Nullable;
 
 /**
@@ -32,17 +34,8 @@ import org.jetbrains.annotations.Nullable;
 public class PyPlainEnterProcessor implements EnterProcessor {
   @Nullable
   private static PyStatementList getStatementList(PsiElement psiElement, Editor editor) {
-    if (psiElement instanceof PyStatementPart) {
-      return ((PyStatementPart)psiElement).getStatementList();
-    }
-    else if (psiElement instanceof PyFunction) {
-      return ((PyFunction)psiElement).getStatementList();
-    }
-    else if (psiElement instanceof PyClass) {
-      return ((PyClass)psiElement).getStatementList();
-    }
-    else if (psiElement instanceof PyWithStatement) {
-      return PsiTreeUtil.getChildOfType(psiElement, PyStatementList.class);
+    if (psiElement instanceof PyStatementListContainer) {
+      return ((PyStatementListContainer)psiElement).getStatementList();
     }
     else {
       final CaretModel caretModel = editor.getCaretModel();