1 files changed, 18 insertions, 15 deletions
diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyClassFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyClassFixer.java
index b9846e1cb894..ce51d6d46824 100644
--- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyClassFixer.java
+++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyClassFixer.java
@@ -17,7 +17,6 @@ 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.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyTokenTypes;
@@ -25,6 +24,9 @@ import com.jetbrains.python.codeInsight.editorActions.smartEnter.PySmartEnterPro
 import com.jetbrains.python.psi.PyArgumentList;
 import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyUtil;
+import org.jetbrains.annotations.NotNull;
+
+import static com.jetbrains.python.psi.PyUtil.sure;
 
 /**
  * Created by IntelliJ IDEA.
@@ -32,21 +34,22 @@ import com.jetbrains.python.psi.PyUtil;
  * Date:   16.04.2010
  * Time:   18:41:08
  */
-public class PyClassFixer implements PyFixer {
-  public void apply(Editor editor, PySmartEnterProcessor processor, PsiElement psiElement) throws IncorrectOperationException {
-    if (psiElement instanceof PyClass) {
-      final PsiElement colon = PyUtil.getChildByFilter(psiElement, TokenSet.create(PyTokenTypes.COLON), 0);
-      if (colon == null) {
-        final PyClass aClass = (PyClass)psiElement;
-        final PyArgumentList argList = PsiTreeUtil.getChildOfType(aClass, PyArgumentList.class);
-        int offset = argList.getTextRange().getEndOffset();
-        String textToInsert = ":";
-        if (aClass.getNameNode() == null) {
-          processor.registerUnresolvedError(argList.getTextRange().getEndOffset() + 1);
-          textToInsert = " :";
-        }
-        editor.getDocument().insertString(offset, textToInsert);
+public class PyClassFixer extends PyFixer<PyClass> {
+  public PyClassFixer() {
+    super(PyClass.class);
+  }
+
+  public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyClass pyClass) throws IncorrectOperationException {
+    final PsiElement colon = PyUtil.getFirstChildOfType(pyClass, PyTokenTypes.COLON);
+    if (colon == null) {
+      final PyArgumentList argList = PsiTreeUtil.getChildOfType(pyClass, PyArgumentList.class);
+      final int offset = sure(argList).getTextRange().getEndOffset();
+      String textToInsert = ":";
+      if (pyClass.getNameNode() == null) {
+        processor.registerUnresolvedError(argList.getTextRange().getEndOffset() + 1);
+        textToInsert = " :";
       }
+      editor.getDocument().insertString(offset, textToInsert);
     }
   }
 }