1 files changed, 22 insertions, 18 deletions
diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyExceptFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyExceptFixer.java
index 8e3532cb49b3..67d5ff5c9ee5 100644
--- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyExceptFixer.java
+++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyExceptFixer.java
@@ -17,13 +17,15 @@ package com.jetbrains.python.codeInsight.editorActions.smartEnter.fixers;
 
 import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.tree.TokenSet;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.codeInsight.editorActions.smartEnter.PySmartEnterProcessor;
 import com.jetbrains.python.psi.PyExceptPart;
 import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.PyUtil;
+import org.jetbrains.annotations.NotNull;
+
+import static com.jetbrains.python.psi.PyUtil.sure;
 
 /**
  * Created by IntelliJ IDEA.
@@ -31,24 +33,26 @@ import com.jetbrains.python.psi.PyUtil;
  * Date:   22.04.2010
  * Time:   18:13:34
  */
-public class PyExceptFixer implements PyFixer {
-  public void apply(Editor editor, PySmartEnterProcessor processor, PsiElement psiElement) throws IncorrectOperationException {
-    if (psiElement instanceof PyExceptPart) {
-      PyExceptPart exceptPart = (PyExceptPart)psiElement;
-      final PsiElement colon = PyUtil.getChildByFilter(exceptPart, TokenSet.create(PyTokenTypes.COLON), 0);
-      if (colon == null) {
-        int offset = PyUtil.getChildByFilter(exceptPart,
-                                             TokenSet.create(PyTokenTypes.EXCEPT_KEYWORD), 0).getTextRange().getEndOffset();
-        final PyExpression exceptClass = exceptPart.getExceptClass();
-        if (exceptClass != null) {
-          offset = exceptClass.getTextRange().getEndOffset();
-        }
-        final PyExpression target = exceptPart.getTarget();
-        if (target != null) {
-          offset = target.getTextRange().getEndOffset();
-        }
-        editor.getDocument().insertString(offset, ":");
+public class PyExceptFixer extends PyFixer<PyExceptPart> {
+  public PyExceptFixer() {
+    super(PyExceptPart.class);
+  }
+
+  @Override
+  public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyExceptPart exceptPart) throws IncorrectOperationException {
+    final PsiElement colon = PyUtil.getFirstChildOfType(exceptPart, PyTokenTypes.COLON);
+    if (colon == null) {
+      final PsiElement exceptToken = PyUtil.getFirstChildOfType(exceptPart, PyTokenTypes.EXCEPT_KEYWORD);
+      int offset = sure(exceptToken).getTextRange().getEndOffset();
+      final PyExpression exceptClass = exceptPart.getExceptClass();
+      if (exceptClass != null) {
+        offset = exceptClass.getTextRange().getEndOffset();
+      }
+      final PyExpression target = exceptPart.getTarget();
+      if (target != null) {
+        offset = target.getTextRange().getEndOffset();
       }
+      editor.getDocument().insertString(offset, ":");
     }
   }
 }