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 a3e189bb9ebd..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; @@ -27,6 +26,8 @@ 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. * Author: Alexey.Ivanov @@ -39,10 +40,10 @@ public class PyClassFixer extends PyFixer { } public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyClass pyClass) throws IncorrectOperationException { - final PsiElement colon = PyUtil.getChildByFilter(pyClass, TokenSet.create(PyTokenTypes.COLON), 0); + final PsiElement colon = PyUtil.getFirstChildOfType(pyClass, PyTokenTypes.COLON); if (colon == null) { final PyArgumentList argList = PsiTreeUtil.getChildOfType(pyClass, PyArgumentList.class); - final int offset = argList.getTextRange().getEndOffset(); + final int offset = sure(argList).getTextRange().getEndOffset(); String textToInsert = ":"; if (pyClass.getNameNode() == null) { processor.registerUnresolvedError(argList.getTextRange().getEndOffset() + 1); diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyConditionalStatementPartFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyConditionalStatementPartFixer.java index c37e210bf8a1..f71c61261103 100644 --- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyConditionalStatementPartFixer.java +++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyConditionalStatementPartFixer.java @@ -27,6 +27,8 @@ 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. * Author: Alexey.Ivanov @@ -43,7 +45,7 @@ public class PyConditionalStatementPartFixer extends PyFixer { @Override public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyExceptPart exceptPart) throws IncorrectOperationException { - final PsiElement colon = PyUtil.getChildByFilter(exceptPart, TokenSet.create(PyTokenTypes.COLON), 0); + final PsiElement colon = PyUtil.getFirstChildOfType(exceptPart, PyTokenTypes.COLON); if (colon == null) { - int offset = PyUtil.getChildByFilter(exceptPart, TokenSet.create(PyTokenTypes.EXCEPT_KEYWORD), 0).getTextRange().getEndOffset(); + 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(); diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyForPartFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyForPartFixer.java index a33e2ba65fab..eefe5cd1fc3e 100644 --- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyForPartFixer.java +++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyForPartFixer.java @@ -18,13 +18,14 @@ package com.jetbrains.python.codeInsight.editorActions.smartEnter.fixers; import com.intellij.openapi.editor.Document; import com.intellij.openapi.editor.Editor; import com.intellij.psi.PsiElement; -import com.intellij.psi.tree.TokenSet; import com.jetbrains.python.PyTokenTypes; import com.jetbrains.python.codeInsight.editorActions.smartEnter.PySmartEnterProcessor; import com.jetbrains.python.psi.PyForPart; import com.jetbrains.python.psi.PyUtil; import org.jetbrains.annotations.NotNull; +import static com.jetbrains.python.psi.PyUtil.sure; + /** * Created by IntelliJ IDEA. * Author: Alexey.Ivanov @@ -38,36 +39,35 @@ public class PyForPartFixer extends PyFixer { @Override public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyForPart forPart) { - final PsiElement colon = PyUtil.getChildByFilter(forPart, TokenSet.create(PyTokenTypes.COLON), 0); - final Document document = editor.getDocument(); - final PsiElement forToken = PyUtil.getChildByFilter(forPart, - TokenSet.create(PyTokenTypes.FOR_KEYWORD), 0); - if (colon == null) { - String textToInsert = ":"; - PsiElement sourceOrTarget = forPart.getSource(); - PsiElement positionToInsert = sourceOrTarget; - if (sourceOrTarget == null) { - sourceOrTarget = forPart.getTarget(); - final PsiElement inToken = PyUtil.getChildByFilter(forPart, TokenSet.create(PyTokenTypes.IN_KEYWORD), 0); - if (inToken == null) { - if (sourceOrTarget == null) { - positionToInsert = forToken; - textToInsert = " in :"; - processor.registerUnresolvedError(positionToInsert.getTextRange().getEndOffset() + 1); - } - else { - positionToInsert = sourceOrTarget; - textToInsert = " in :"; - processor.registerUnresolvedError(positionToInsert.getTextRange().getEndOffset() + 4); - } - } - else { - positionToInsert = inToken; - textToInsert = " :"; + final PsiElement colon = PyUtil.getFirstChildOfType(forPart, PyTokenTypes.COLON); + final Document document = editor.getDocument(); + final PsiElement forToken = PyUtil.getFirstChildOfType(forPart, PyTokenTypes.FOR_KEYWORD); + if (colon == null) { + String textToInsert = ":"; + PsiElement sourceOrTarget = forPart.getSource(); + PsiElement positionToInsert = sourceOrTarget; + if (sourceOrTarget == null) { + sourceOrTarget = forPart.getTarget(); + final PsiElement inToken = PyUtil.getFirstChildOfType(forPart, PyTokenTypes.IN_KEYWORD); + if (inToken == null) { + if (sourceOrTarget == null) { + positionToInsert = sure(forToken); + textToInsert = " in :"; processor.registerUnresolvedError(positionToInsert.getTextRange().getEndOffset() + 1); } + else { + positionToInsert = sourceOrTarget; + textToInsert = " in :"; + processor.registerUnresolvedError(positionToInsert.getTextRange().getEndOffset() + 4); + } + } + else { + positionToInsert = inToken; + textToInsert = " :"; + processor.registerUnresolvedError(positionToInsert.getTextRange().getEndOffset() + 1); } - document.insertString(positionToInsert.getTextRange().getEndOffset(), textToInsert); } + document.insertString(positionToInsert.getTextRange().getEndOffset(), textToInsert); } + } } diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyFunctionFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyFunctionFixer.java index 0765aa2a16f2..ad959f80ccae 100644 --- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyFunctionFixer.java +++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyFunctionFixer.java @@ -18,7 +18,6 @@ package com.jetbrains.python.codeInsight.editorActions.smartEnter.fixers; import com.intellij.openapi.editor.Document; 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; @@ -41,7 +40,7 @@ public class PyFunctionFixer extends PyFixer { @Override public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyFunction function) throws IncorrectOperationException { - final PsiElement colon = PyUtil.getChildByFilter(function, TokenSet.create(PyTokenTypes.COLON), 0); + final PsiElement colon = PyUtil.getFirstChildOfType(function, PyTokenTypes.COLON); if (colon == null) { final PyParameterList parameterList = function.getParameterList(); final Document document = editor.getDocument(); diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyMissingBracesFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyMissingBracesFixer.java index ea3fb3162d89..eb97d2bdb372 100644 --- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyMissingBracesFixer.java +++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyMissingBracesFixer.java @@ -34,9 +34,10 @@ public class PyMissingBracesFixer extends PyFixer { } @Override - public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyElement psiElement) throws IncorrectOperationException { + public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyElement psiElement) + throws IncorrectOperationException { if (psiElement instanceof PySetLiteralExpression || psiElement instanceof PyDictLiteralExpression) { - PsiElement lastChild = PyUtil.getFirstNonCommentBefore(psiElement.getLastChild()); + final PsiElement lastChild = PyUtil.getFirstNonCommentBefore(psiElement.getLastChild()); if (lastChild != null && !"}".equals(lastChild.getText())) { editor.getDocument().insertString(lastChild.getTextRange().getEndOffset(), "}"); } @@ -44,7 +45,7 @@ public class PyMissingBracesFixer extends PyFixer { else if (psiElement instanceof PyListLiteralExpression || psiElement instanceof PySliceExpression || psiElement instanceof PySubscriptionExpression) { - PsiElement lastChild = PyUtil.getFirstNonCommentBefore(psiElement.getLastChild()); + final PsiElement lastChild = PyUtil.getFirstNonCommentBefore(psiElement.getLastChild()); if (lastChild != null && !"]".equals(lastChild.getText())) { editor.getDocument().insertString(lastChild.getTextRange().getEndOffset(), "]"); } diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyUnconditionalStatementPartFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyUnconditionalStatementPartFixer.java index 8105723893d7..6aace3972a1b 100644 --- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyUnconditionalStatementPartFixer.java +++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyUnconditionalStatementPartFixer.java @@ -24,6 +24,8 @@ import com.jetbrains.python.codeInsight.editorActions.smartEnter.PySmartEnterPro import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; +import static com.jetbrains.python.psi.PyUtil.sure; + /** * Created by IntelliJ IDEA. * Author: Alexey.Ivanov @@ -39,11 +41,11 @@ public class PyUnconditionalStatementPartFixer extends PyFixer { public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyElement psiElement) throws IncorrectOperationException { if (PyUtil.instanceOf(psiElement, PyElsePart.class, PyTryPart.class, PyFinallyPart.class)) { - final PsiElement colon = PyUtil.getChildByFilter(psiElement, TokenSet.create(PyTokenTypes.COLON), 0); + final PsiElement colon = PyUtil.getFirstChildOfType(psiElement, PyTokenTypes.COLON); if (colon == null) { final TokenSet keywords = TokenSet.create(PyTokenTypes.ELSE_KEYWORD, PyTokenTypes.TRY_KEYWORD, PyTokenTypes.FINALLY_KEYWORD); final PsiElement keywordToken = PyUtil.getChildByFilter(psiElement, keywords, 0); - editor.getDocument().insertString(keywordToken.getTextRange().getEndOffset(), ":"); + editor.getDocument().insertString(sure(keywordToken).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 index b22ceed996a4..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,19 +15,17 @@ */ 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; @@ -41,8 +39,8 @@ public class PyWithFixer extends PyFixer { @Override public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyWithStatement withStatement) throws IncorrectOperationException { - final PsiElement colonToken = getFirstChildOfType(withStatement, PyTokenTypes.COLON); - final PsiElement withToken = getFirstChildOfType(withStatement, PyTokenTypes.WITH_KEYWORD); + 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(); @@ -56,7 +54,7 @@ public class PyWithFixer extends PyFixer { else { final PyExpression expression = lastItem.getExpression(); insertAt = expression.getTextRange().getEndOffset(); - final PsiElement asToken = getFirstChildOfType(lastItem, PyTokenTypes.AS_KEYWORD); + final PsiElement asToken = PyUtil.getFirstChildOfType(lastItem, PyTokenTypes.AS_KEYWORD); if (asToken != null) { insertAt = asToken.getTextRange().getEndOffset(); final PyExpression target = lastItem.getTarget(); @@ -72,10 +70,4 @@ public class PyWithFixer extends PyFixer { 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/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index f1dd16078696..cd4e582bb6f2 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/python/src/com/jetbrains/python/psi/PyUtil.java @@ -900,6 +900,20 @@ public class PyUtil { return null; } + /** + * Returns first child psi element with specified element type or {@code null} if no such element exists. + * Semantically it's the same as {@code getChildByFilter(element, TokenSet.create(type), 0)}. + * + * @param element tree parent node + * @param type element type expected + * @return child element described + */ + @Nullable + public static PsiElement getFirstChildOfType(@NotNull final PsiElement element, @NotNull PyElementType type) { + final ASTNode child = element.getNode().findChildByType(type); + return child != null ? child.getPsi() : null; + } + /** * If argument is a PsiDirectory, turn it into a PsiFile that points to __init__.py in that directory. * If there's no __init__.py there, null is returned, there's no point to resolve to a dir which is not a package.