diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/PySmartEnterProcessor.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/PySmartEnterProcessor.java index aaedcf60eae4..2da1b257a9a0 100644 --- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/PySmartEnterProcessor.java +++ b/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 ourProcessors = ImmutableList.of(new PyCommentBreakerEnterProcessor(), new PyPlainEnterProcessor()); diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/enterProcessors/PyPlainEnterProcessor.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/enterProcessors/PyPlainEnterProcessor.java index d6717905a034..ac2e31252dcf 100644 --- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/enterProcessors/PyPlainEnterProcessor.java +++ b/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()); 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 new file mode 100644 index 000000000000..b6d099ef7cb5 --- /dev/null +++ b/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; + } +} diff --git a/python/testData/codeInsight/smartEnter/withTargetIncomplete.py b/python/testData/codeInsight/smartEnter/withTargetIncomplete.py new file mode 100644 index 000000000000..a3214a35ffe1 --- /dev/null +++ b/python/testData/codeInsight/smartEnter/withTargetIncomplete.py @@ -0,0 +1 @@ +with open('file.txt') as \ No newline at end of file diff --git a/python/testData/codeInsight/smartEnter/withTargetIncomplete_after.py b/python/testData/codeInsight/smartEnter/withTargetIncomplete_after.py new file mode 100644 index 000000000000..72892c552212 --- /dev/null +++ b/python/testData/codeInsight/smartEnter/withTargetIncomplete_after.py @@ -0,0 +1 @@ +with open('file.txt') as : \ No newline at end of file diff --git a/python/testData/codeInsight/smartEnter/withTargetOmitted.py b/python/testData/codeInsight/smartEnter/withTargetOmitted.py new file mode 100644 index 000000000000..c1ba50b34f80 --- /dev/null +++ b/python/testData/codeInsight/smartEnter/withTargetOmitted.py @@ -0,0 +1 @@ +with open('file.txt') \ No newline at end of file diff --git a/python/testData/codeInsight/smartEnter/withTargetOmitted_after.py b/python/testData/codeInsight/smartEnter/withTargetOmitted_after.py new file mode 100644 index 000000000000..a9e851f0023d --- /dev/null +++ b/python/testData/codeInsight/smartEnter/withTargetOmitted_after.py @@ -0,0 +1,2 @@ +with open('file.txt'): + \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PySmartEnterTest.java b/python/testSrc/com/jetbrains/python/PySmartEnterTest.java index 959b1f752ee8..b8421c9d7e18 100644 --- a/python/testSrc/com/jetbrains/python/PySmartEnterTest.java +++ b/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(); + } }