1 files changed, 34 insertions, 40 deletions

diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyWithFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyWithFixer.java
index b6d099ef7cb5..ec236d9242f5 100644
--- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyWithFixer.java
+++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyWithFixer.java
@@ -15,65 +15,59 @@
  */
 package com.jetbrains.python.codeInsight.editorActions.smartEnter.fixers;
 
-import com.intellij.lang.ASTNode;
 import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiElement;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.codeInsight.editorActions.smartEnter.PySmartEnterProcessor;
-import com.jetbrains.python.psi.PyElementType;
 import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.PyUtil;
 import com.jetbrains.python.psi.PyWithItem;
 import com.jetbrains.python.psi.PyWithStatement;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 
 import static com.jetbrains.python.psi.PyUtil.sure;
 
 /**
  * @author Mikhail Golubev
  */
-public class PyWithFixer implements PyFixer {
-  public void apply(Editor editor, PySmartEnterProcessor processor, PsiElement psiElement) throws IncorrectOperationException {
-    if (psiElement instanceof PyWithStatement) {
-      final PyWithStatement withStatement = (PyWithStatement)psiElement;
-      final PsiElement colonToken = getFirstChildOfType(psiElement, PyTokenTypes.COLON);
-      final PsiElement withToken = getFirstChildOfType(withStatement, PyTokenTypes.WITH_KEYWORD);
-      final Document document = editor.getDocument();
-      if (colonToken == null) {
-        int insertAt = sure(withToken).getTextRange().getEndOffset();
-        String textToInsert = ":";
-        final PyWithItem[] withItems = withStatement.getWithItems();
-        final PyWithItem lastItem = withItems.length != 0 ? withItems[withItems.length - 1] : null;
-        if (lastItem == null || lastItem.getExpression() == null) {
-          textToInsert = " :";
-          processor.registerUnresolvedError(insertAt + 1);
-        }
-        else {
-          final PyExpression expression = lastItem.getExpression();
-          insertAt = expression.getTextRange().getEndOffset();
-          final PsiElement asToken = getFirstChildOfType(lastItem, PyTokenTypes.AS_KEYWORD);
-          if (asToken != null) {
-            insertAt = asToken.getTextRange().getEndOffset();
-            final PyExpression target = lastItem.getTarget();
-            if (target != null) {
-              insertAt = target.getTextRange().getEndOffset();
-            }
-            else {
-              textToInsert = " :";
-              processor.registerUnresolvedError(insertAt + 1);
-            }
+public class PyWithFixer extends PyFixer<PyWithStatement> {
+  public PyWithFixer() {
+    super(PyWithStatement.class);
+  }
+
+  @Override
+  public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyWithStatement withStatement) throws IncorrectOperationException {
+    final PsiElement colonToken = PyUtil.getFirstChildOfType(withStatement, PyTokenTypes.COLON);
+    final PsiElement withToken = PyUtil.getFirstChildOfType(withStatement, PyTokenTypes.WITH_KEYWORD);
+    final Document document = editor.getDocument();
+    if (colonToken == null) {
+      int insertAt = sure(withToken).getTextRange().getEndOffset();
+      String textToInsert = ":";
+      final PyWithItem[] withItems = withStatement.getWithItems();
+      final PyWithItem lastItem = withItems.length != 0 ? withItems[withItems.length - 1] : null;
+      if (lastItem == null || lastItem.getExpression() == null) {
+        textToInsert = " :";
+        processor.registerUnresolvedError(insertAt + 1);
+      }
+      else {
+        final PyExpression expression = lastItem.getExpression();
+        insertAt = expression.getTextRange().getEndOffset();
+        final PsiElement asToken = PyUtil.getFirstChildOfType(lastItem, PyTokenTypes.AS_KEYWORD);
+        if (asToken != null) {
+          insertAt = asToken.getTextRange().getEndOffset();
+          final PyExpression target = lastItem.getTarget();
+          if (target != null) {
+            insertAt = target.getTextRange().getEndOffset();
+          }
+          else {
+            textToInsert = " :";
+            processor.registerUnresolvedError(insertAt + 1);
           }
         }
-        document.insertString(insertAt, textToInsert);
       }
+      document.insertString(insertAt, textToInsert);
     }
   }
-
-  @Nullable
-  private static PsiElement getFirstChildOfType(@NotNull final PsiElement element, @NotNull PyElementType type) {
-    final ASTNode child = element.getNode().findChildByType(type);
-    return child != null ? child.getPsi() : null;
-  }
 }