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