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