PY-12877 Complete statement doesn't insert colon for 'with' statement

python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/PySmartEnterProcessor.java

@@ -61,6 +61,7 @@ public class PySmartEnterProcessor extends SmartEnterProcessor {
     .add(new PyParameterListFixer())
     .add(new PyFunctionFixer())
     .add(new PyClassFixer())
+    .add(new PyWithFixer())
     .build();
   private static final List<EnterProcessor> ourProcessors = ImmutableList.of(new PyCommentBreakerEnterProcessor(),
                                                                              new PyPlainEnterProcessor());

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

@@ -20,10 +20,7 @@ import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.codeInsight.editorActions.smartEnter.SmartEnterUtil;
-import com.jetbrains.python.psi.PyClass;
-import com.jetbrains.python.psi.PyFunction;
-import com.jetbrains.python.psi.PyStatementList;
-import com.jetbrains.python.psi.PyStatementPart;
+import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.Nullable;
 
 /**
@@ -43,10 +40,14 @@ public class PyPlainEnterProcessor implements EnterProcessor {
     }
     else if (psiElement instanceof PyClass) {
       return ((PyClass)psiElement).getStatementList();
-    } else {
-    final CaretModel caretModel = editor.getCaretModel();
-    final PsiElement atCaret = psiElement.getContainingFile().findElementAt(caretModel.getOffset());
-      PyStatementPart statementPart = PsiTreeUtil.getParentOfType(atCaret, PyStatementPart.class);
+    }
+    else if (psiElement instanceof PyWithStatement) {
+      return PsiTreeUtil.getChildOfType(psiElement, PyStatementList.class);
+    }
+    else {
+      final CaretModel caretModel = editor.getCaretModel();
+      final PsiElement atCaret = psiElement.getContainingFile().findElementAt(caretModel.getOffset());
+      final PyStatementPart statementPart = PsiTreeUtil.getParentOfType(atCaret, PyStatementPart.class);
       if (statementPart != null) {
         return statementPart.getStatementList();
       }
@@ -55,7 +56,7 @@ public class PyPlainEnterProcessor implements EnterProcessor {
   }
 
   public boolean doEnter(Editor editor, PsiElement psiElement, boolean isModified) {
-    PyStatementList statementList = getStatementList(psiElement, editor);
+    final PyStatementList statementList = getStatementList(psiElement, editor);
     if (statementList != null && statementList.getStatements().length == 0) {
       SmartEnterUtil.plainEnter(editor);
       //editor.getCaretModel().moveToOffset(statementList.getTextRange().getEndOffset());

python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyWithFixer.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2000-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+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.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);
+            }
+          }
+        }
+        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;
+  }
+}

python/testData/codeInsight/smartEnter/withTargetIncomplete.py

@@ -0,0 +1 @@
+with open('file.txt') <caret>as

python/testData/codeInsight/smartEnter/withTargetIncomplete_after.py

@@ -0,0 +1 @@
+with open('file.txt') as <caret>:

python/testData/codeInsight/smartEnter/withTargetOmitted.py

@@ -0,0 +1 @@
+with o<caret>pen('file.txt')

python/testData/codeInsight/smartEnter/withTargetOmitted_after.py

@@ -0,0 +1,2 @@
+with open('file.txt'):
+    <caret>

python/testSrc/com/jetbrains/python/PySmartEnterTest.java

@@ -169,4 +169,14 @@ public class PySmartEnterTest extends PyTestCase {
       pyCodeInsightSettings.INSERT_TYPE_DOCSTUB = oldInsertType;
     }
   }
+
+  // PY-12877
+  public void testWithTargetOmitted() {
+    doTest();
+  }
+
+  // PY-12877
+  public void testWithTargetIncomplete() {
+    doTest();
+  }
 }