Examples with PyWithItem com.jetbrains.python.psi.PyWithItem used on opensource projects

Example 1 with PyWithItem

use of com.jetbrains.python.psi.PyWithItem in project intellij-community by JetBrains.

the class PyWithFixer method doApply.

@Override
public void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PyWithStatement withStatement) throws IncorrectOperationException {
    final PsiElement colonToken = PyPsiUtils.getFirstChildOfType(withStatement, PyTokenTypes.COLON);
    final PsiElement withToken = PyPsiUtils.getFirstChildOfType(withStatement, PyTokenTypes.WITH_KEYWORD);
    final Document document = editor.getDocument();
    if (colonToken == null && withToken != null) {
        int insertAt = withToken.getTextRange().getEndOffset();
        String textToInsert = ":";
        final PyWithItem lastItem = ArrayUtil.getLastElement(withStatement.getWithItems());
        if (lastItem == null || lastItem.getExpression() == null) {
            textToInsert = " :";
            processor.registerUnresolvedError(insertAt + 1);
        } else {
            final PyExpression expression = lastItem.getExpression();
            insertAt = expression.getTextRange().getEndOffset();
            final PsiElement asToken = PyPsiUtils.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);
    }
}