@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 && 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 = 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);
}
}
