@Override
public boolean isAsyncAllowed() {
final LanguageLevel languageLevel = LanguageLevel.forElement(this);
final String functionName = getName();
return languageLevel.isAtLeast(LanguageLevel.PYTHON35)
&& (functionName == null
|| ArrayUtil.contains(
functionName, PyNames.AITER, PyNames.ANEXT, PyNames.AENTER, PyNames.AEXIT)
|| !PyNames.getBuiltinMethods(languageLevel).containsKey(functionName));
}
private void addToExistingImport(PyImportElement src) {
final PyElementGenerator gen = PyElementGenerator.getInstance(myTarget.getProject());
// did user choose 'import' or 'from import'?
PsiElement parent = src.getParent();
if (parent instanceof PyFromImportStatement) {
// add another import element right after the one we got
PsiElement newImportElement = gen.createImportElement(LanguageLevel.getDefault(), myName);
parent.add(newImportElement);
} else { // just 'import'
// all we need is to qualify our target
myTarget.replace(
gen.createExpressionFromText(
LanguageLevel.forElement(myTarget), src.getVisibleName() + "." + myName));
}
}
@Nullable
@Override
public PyType getReturnType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final PyType returnType = typeProvider.getReturnType(this, context);
if (returnType != null) {
returnType.assertValid(typeProvider.toString());
return returnType;
}
}
if (context.maySwitchToAST(this)
&& LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON30)) {
PyAnnotation anno = getAnnotation();
if (anno != null) {
PyClass pyClass = anno.resolveToClass();
if (pyClass != null) {
return new PyClassTypeImpl(pyClass, false);
}
}
}
final PyType docStringType = getReturnTypeFromDocString();
if (docStringType != null) {
docStringType.assertValid("from docstring");
return docStringType;
}
if (context.allowReturnTypes(this)) {
final Ref<? extends PyType> yieldTypeRef = getYieldStatementType(context);
if (yieldTypeRef != null) {
return yieldTypeRef.get();
}
return getReturnStatementType(context);
}
return null;
}
@Override
protected void moveOffsetAfter(boolean success) {
if (success && (myPanel != null && myPanel.getInitPlace() != InitPlace.SAME_METHOD)
|| myOperation.getInplaceInitPlace() != InitPlace.SAME_METHOD) {
final AccessToken accessToken =
ApplicationManager.getApplication().acquireWriteActionLock(getClass());
try {
final PyAssignmentStatement initializer =
PsiTreeUtil.getParentOfType(myTarget, PyAssignmentStatement.class);
assert initializer != null;
final Function<String, PyStatement> callback =
FunctionUtil.<String, PyStatement>constant(initializer);
final PyClass pyClass = PyUtil.getContainingClassOrSelf(initializer);
InitPlace initPlace =
myPanel != null ? myPanel.getInitPlace() : myOperation.getInplaceInitPlace();
if (initPlace == InitPlace.CONSTRUCTOR) {
AddFieldQuickFix.addFieldToInit(myProject, pyClass, "", callback);
} else if (initPlace == InitPlace.SET_UP) {
addFieldToSetUp(pyClass, callback);
}
if (myOperation.getOccurrences().size() > 0) {
initializer.delete();
} else {
final PyExpression copy =
PyElementGenerator.getInstance(myProject)
.createExpressionFromText(
LanguageLevel.forElement(myTarget), myTarget.getText());
initializer.replace(copy);
}
initializer.delete();
} finally {
accessToken.finish();
}
}
}
@Nullable
private PyType getGenericReturnType(
@NotNull TypeEvalContext typeEvalContext, @Nullable PyQualifiedExpression callSite) {
if (typeEvalContext.maySwitchToAST(this)
&& LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON30)) {
PyAnnotation anno = getAnnotation();
if (anno != null) {
PyClass pyClass = anno.resolveToClass();
if (pyClass != null) {
return new PyClassTypeImpl(pyClass, false);
}
}
}
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final PyType returnType = typeProvider.getReturnType(this, callSite, typeEvalContext);
if (returnType != null) {
returnType.assertValid(typeProvider.toString());
return returnType;
}
}
final PyType docStringType = getReturnTypeFromDocString();
if (docStringType != null) {
docStringType.assertValid("from docstring");
return docStringType;
}
if (typeEvalContext.allowReturnTypes(this)) {
final PyType yieldType = getYieldStatementType(typeEvalContext);
if (yieldType != null) {
return yieldType;
}
return getReturnStatementType(typeEvalContext);
}
return null;
}
@Override
protected PyAssignmentStatement createDeclaration(
Project project, String assignmentText, PsiElement anchor) {
final PyFunction container = PsiTreeUtil.getParentOfType(anchor, PyFunction.class);
String selfName = PyUtil.getFirstParameterName(container);
final LanguageLevel langLevel = LanguageLevel.forElement(anchor);
return PyElementGenerator.getInstance(project)
.createFromText(langLevel, PyAssignmentStatement.class, selfName + "." + assignmentText);
}
public PyStatement fun(String self_name) {
if (PyNames.CANONICAL_SELF.equals(self_name)) {
return (PyStatement) myDeclaration;
}
final String text = myDeclaration.getText();
final Project project = myDeclaration.getProject();
return PyElementGenerator.getInstance(project)
.createFromText(
LanguageLevel.getDefault(),
PyStatement.class,
text.replaceFirst(PyNames.CANONICAL_SELF + "\\.", self_name + "."));
}
@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);
}
private void addImportStatement(ImportCandidateHolder item) {
final Project project = myTarget.getProject();
final PyElementGenerator gen = PyElementGenerator.getInstance(project);
AddImportHelper.ImportPriority priority =
AddImportHelper.getImportPriority(myTarget, item.getFile());
PsiFile file = myTarget.getContainingFile();
InjectedLanguageManager manager = InjectedLanguageManager.getInstance(project);
if (manager.isInjectedFragment(file)) {
file = manager.getTopLevelFile(myTarget);
}
// We are trying to import top-level module or package which thus cannot be qualified
if (isRoot(item.getFile())) {
if (myImportLocally) {
AddImportHelper.addLocalImportStatement(myTarget, myName);
} else {
AddImportHelper.addImportStatement(file, myName, item.getAsName(), priority, null);
}
} else {
final QualifiedName path = item.getPath();
final String qualifiedName = path != null ? path.toString() : "";
if (myUseQualifiedImport) {
String nameToImport = qualifiedName;
if (item.getImportable() instanceof PsiFileSystemItem) {
nameToImport += "." + myName;
}
if (myImportLocally) {
AddImportHelper.addLocalImportStatement(myTarget, nameToImport);
} else {
AddImportHelper.addImportStatement(file, nameToImport, item.getAsName(), priority, null);
}
myTarget.replace(
gen.createExpressionFromText(
LanguageLevel.forElement(myTarget), qualifiedName + "." + myName));
} else {
if (myImportLocally) {
AddImportHelper.addLocalFromImportStatement(myTarget, qualifiedName, myName);
} else {
AddImportHelper.addFromImportStatement(
file, qualifiedName, myName, item.getAsName(), priority, null);
}
}
}
}
