@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
public String extractDeprecationMessage() {
PyStatementList statementList = getStatementList();
if (statementList == null) {
return null;
}
return extractDeprecationMessage(Arrays.asList(statementList.getStatements()));
}
public boolean doEnter(Editor editor, PsiElement psiElement, boolean isModified) {
final PyStatementList statementList = getStatementList(psiElement, editor);
if (statementList != null && statementList.getStatements().length == 0) {
SmartEnterUtil.plainEnter(editor);
// editor.getCaretModel().moveToOffset(statementList.getTextRange().getEndOffset());
return true;
}
return false;
}
@Nullable
public PyType getReturnStatementType(TypeEvalContext typeEvalContext) {
final ReturnVisitor visitor = new ReturnVisitor(this, typeEvalContext);
final PyStatementList statements = getStatementList();
statements.accept(visitor);
if (isGeneratedStub() && !visitor.myHasReturns) {
if (PyNames.INIT.equals(getName())) {
return PyNoneType.INSTANCE;
}
return null;
}
return visitor.result();
}
@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
private Ref<? extends PyType> getYieldStatementType(@NotNull final TypeEvalContext context) {
Ref<PyType> elementType = null;
final PyBuiltinCache cache = PyBuiltinCache.getInstance(this);
final PyStatementList statements = getStatementList();
final Set<PyType> types = new LinkedHashSet<>();
statements.accept(
new PyRecursiveElementVisitor() {
@Override
public void visitPyYieldExpression(PyYieldExpression node) {
final PyExpression expr = node.getExpression();
final PyType type = expr != null ? context.getType(expr) : null;
if (node.isDelegating() && type instanceof PyCollectionType) {
final PyCollectionType collectionType = (PyCollectionType) type;
// TODO: Select the parameter types that matches T in Iterable[T]
final List<PyType> elementTypes = collectionType.getElementTypes(context);
types.add(elementTypes.isEmpty() ? null : elementTypes.get(0));
} else {
types.add(type);
}
}
@Override
public void visitPyFunction(PyFunction node) {
// Ignore nested functions
}
});
final int n = types.size();
if (n == 1) {
elementType = Ref.create(types.iterator().next());
} else if (n > 0) {
elementType = Ref.create(PyUnionType.union(types));
}
if (elementType != null) {
final PyClass generator = cache.getClass(PyNames.FAKE_GENERATOR);
if (generator != null) {
final List<PyType> parameters =
Arrays.asList(elementType.get(), null, getReturnStatementType(context));
return Ref.create(new PyCollectionTypeImpl(generator, false, parameters));
}
}
if (!types.isEmpty()) {
return Ref.create(null);
}
return null;
}
@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;
}
@Nullable
private Ref<? extends PyType> getYieldStatementType(@NotNull final TypeEvalContext context) {
Ref<PyType> elementType = null;
final PyBuiltinCache cache = PyBuiltinCache.getInstance(this);
final PyStatementList statements = getStatementList();
final Set<PyType> types = new LinkedHashSet<PyType>();
if (statements != null) {
statements.accept(
new PyRecursiveElementVisitor() {
@Override
public void visitPyYieldExpression(PyYieldExpression node) {
final PyType type = context.getType(node);
if (node.isDelegating() && type instanceof PyCollectionType) {
final PyCollectionType collectionType = (PyCollectionType) type;
types.add(collectionType.getElementType(context));
} else {
types.add(type);
}
}
@Override
public void visitPyFunction(PyFunction node) {
// Ignore nested functions
}
});
final int n = types.size();
if (n == 1) {
elementType = Ref.create(types.iterator().next());
} else if (n > 0) {
elementType = Ref.create(PyUnionType.union(types));
}
}
if (elementType != null) {
final PyClass generator = cache.getClass(PyNames.FAKE_GENERATOR);
if (generator != null) {
return Ref.create(new PyCollectionTypeImpl(generator, false, elementType.get()));
}
}
if (!types.isEmpty()) {
return Ref.create(null);
}
return null;
}
protected TextRange getResultRange(PyTryExceptStatement tryStatement) {
final PyExceptPart part = tryStatement.getExceptParts()[0];
final PyStatementList list = part.getStatementList();
return list.getTextRange();
}