PyWithStatementImpl implements PyStatementListContainer

python/psi-api/src/com/jetbrains/python/psi/PyWithStatement.java

@@ -18,6 +18,6 @@ package com.jetbrains.python.psi;
 /**
  * @author yole
  */
-public interface PyWithStatement extends PyStatement, NameDefiner {
+public interface PyWithStatement extends PyStatement, NameDefiner, PyStatementListContainer {
   PyWithItem[] getWithItems();
 }

python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/enterProcessors/PyPlainEnterProcessor.java

@@ -35,14 +35,8 @@ public class PyPlainEnterProcessor implements EnterProcessor {
     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);
+    else if (psiElement instanceof PyStatementListContainer) {
+      return ((PyStatementListContainer)psiElement).getStatementList();
     }
     else {
       final CaretModel caretModel = editor.getCaretModel();

python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java

@@ -63,6 +63,14 @@ public class PyWithStatementImpl extends PyElementImpl implements PyWithStatemen
   }
 
   public PyWithItem[] getWithItems() {
-    return childrenToPsi(WITH_ITEM, PyWithItem.EMPTY_ARRAY); 
+    return childrenToPsi(WITH_ITEM, PyWithItem.EMPTY_ARRAY);
+  }
+
+  @Override
+  @NotNull
+  public PyStatementList getStatementList() {
+    final PyStatementList statementList = childToPsi(PyElementTypes.STATEMENT_LIST);
+    assert statementList != null : "Statement list missing for with statement " + getText();
+    return statementList;
   }
 }