@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;
}
@Nullable
private static PsiElement addFieldToSetUp(
PyClass clazz, final Function<String, PyStatement> callback) {
final PyFunction init =
clazz.findMethodByName(PythonUnitTestUtil.TESTCASE_SETUP_NAME, false, null);
if (init != null) {
return AddFieldQuickFix.appendToMethod(init, callback);
}
final PyFunctionBuilder builder =
new PyFunctionBuilder(PythonUnitTestUtil.TESTCASE_SETUP_NAME, clazz);
builder.parameter(PyNames.CANONICAL_SELF);
PyFunction setUp = builder.buildFunction(clazz.getProject(), LanguageLevel.getDefault());
final PyStatementList statements = clazz.getStatementList();
final PsiElement anchor = statements.getFirstChild();
setUp = (PyFunction) statements.addBefore(setUp, anchor);
return AddFieldQuickFix.appendToMethod(setUp, callback);
}
@Nullable
@Override
public PsiComment getTypeComment() {
final PsiComment inlineComment = PyUtil.getCommentOnHeaderLine(this);
if (inlineComment != null
&& PyTypingTypeProvider.getTypeCommentValue(inlineComment.getText()) != null) {
return inlineComment;
}
final PyStatementList statements = getStatementList();
if (statements.getStatements().length != 0) {
final PsiComment comment = as(statements.getFirstChild(), PsiComment.class);
if (comment != null && PyTypingTypeProvider.getTypeCommentValue(comment.getText()) != null) {
return comment;
}
}
return null;
}
