1 files changed, 6 insertions, 8 deletions
diff --git a/python/src/com/jetbrains/python/psi/impl/PyStatementPartImpl.java b/python/src/com/jetbrains/python/psi/impl/PyStatementPartImpl.java
index 46e9a785a9d4..7454d6741382 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyStatementPartImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyStatementPartImpl.java
@@ -15,10 +15,11 @@
  */
 package com.jetbrains.python.psi.impl;
 
-import com.jetbrains.python.psi.PyStatementPart;
-import com.jetbrains.python.psi.PyStatementList;
-import com.jetbrains.python.PyElementTypes;
 import com.intellij.lang.ASTNode;
+import com.jetbrains.python.PyElementTypes;
+import com.jetbrains.python.psi.PyStatementList;
+import com.jetbrains.python.psi.PyStatementPart;
+import org.jetbrains.annotations.NotNull;
 
 /**
  * Abstract statement part implementation; extracts the statements list.
@@ -30,11 +31,8 @@ public abstract class PyStatementPartImpl extends PyElementImpl implements PySta
     super(astNode);
   }
 
+  @NotNull
   public PyStatementList getStatementList() {
-    ASTNode n = getNode().findChildByType(PyElementTypes.STATEMENT_LISTS);
-    if (n != null) {
-      return (PyStatementList)n.getPsi();
-    }
-    return null;
+    return childToPsiNotNull(PyElementTypes.STATEMENT_LIST);
   }
 }