@Override
@Nullable
protected TextRange surroundStatement(
@NotNull Project project, @NotNull Editor editor, @NotNull PsiElement[] elements)
throws IncorrectOperationException {
PyTryExceptStatement tryStatement =
PyElementGenerator.getInstance(project)
.createFromText(LanguageLevel.getDefault(), PyTryExceptStatement.class, getTemplate());
final PsiElement parent = elements[0].getParent();
final PyStatementList statementList = tryStatement.getTryPart().getStatementList();
statementList.addRange(elements[0], elements[elements.length - 1]);
statementList.getFirstChild().delete();
tryStatement = (PyTryExceptStatement) parent.addBefore(tryStatement, elements[0]);
parent.deleteChildRange(elements[0], elements[elements.length - 1]);
final PsiFile psiFile = parent.getContainingFile();
final Document document = psiFile.getViewProvider().getDocument();
final TextRange range = tryStatement.getTextRange();
assert document != null;
final RangeMarker rangeMarker = document.createRangeMarker(range);
final PsiElement element = psiFile.findElementAt(rangeMarker.getStartOffset());
tryStatement = PsiTreeUtil.getParentOfType(element, PyTryExceptStatement.class);
if (tryStatement != null) {
return getResultRange(tryStatement);
}
return null;
}
protected TextRange getResultRange(PyTryExceptStatement tryStatement) {
final PyExceptPart part = tryStatement.getExceptParts()[0];
final PyStatementList list = part.getStatementList();
return list.getTextRange();
}