Examples with PyStatementSurrounder com.jetbrains.python.refactoring.surround.surrounders.statements.PyStatementSurrounder used on opensource projects

Example 1 with PyStatementSurrounder

use of com.jetbrains.python.refactoring.surround.surrounders.statements.PyStatementSurrounder in project intellij-community by JetBrains.

the class PyMainPostfixTemplate method getSurrounder.

@NotNull
@Override
protected Surrounder getSurrounder() {
    return new PyStatementSurrounder() {

        @Nullable
        @Override
        protected TextRange surroundStatement(@NotNull Project project, @NotNull Editor editor, @NotNull PsiElement[] elements) throws IncorrectOperationException {
            PyIfStatement ifStatement = PyElementGenerator.getInstance(project).createFromText(LanguageLevel.forElement(elements[0]), PyIfStatement.class, "if __name__ == '__main__':\n expr");
            ifStatement = (PyIfStatement) CodeStyleManager.getInstance(project).reformat(ifStatement);
            final PsiElement parent = elements[0].getParent();
            ifStatement = (PyIfStatement) parent.addBefore(ifStatement, elements[0]);
            final PyStatementList statementList = ifStatement.getIfPart().getStatementList();
            statementList.addRange(elements[0], elements[elements.length - 1]);
            statementList.getFirstChild().delete();
            parent.deleteChildRange(elements[0], elements[elements.length - 1]);
            return TextRange.from(statementList.getTextRange().getEndOffset(), 0);
        }

        @Override
        public String getTemplateDescription() {
            return DESCR;
        }
    };
}