cleanup
[idea/community.git] / python / src / com / jetbrains / python / codeInsight / postfix / PyExpressionAsConditionSurrounder.java
1 /*
2  * Copyright 2000-2016 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.jetbrains.python.codeInsight.postfix;
17
18 import com.intellij.openapi.editor.Editor;
19 import com.intellij.openapi.project.Project;
20 import com.intellij.openapi.util.TextRange;
21 import com.intellij.psi.PsiDocumentManager;
22 import com.intellij.psi.codeStyle.CodeStyleManager;
23 import com.intellij.util.IncorrectOperationException;
24 import com.jetbrains.python.psi.*;
25 import com.jetbrains.python.refactoring.surround.surrounders.expressions.PyExpressionSurrounder;
26 import org.jetbrains.annotations.NotNull;
27 import org.jetbrains.annotations.Nullable;
28
29 public abstract class PyExpressionAsConditionSurrounder extends PyExpressionSurrounder {
30
31   protected abstract String getTextToGenerate();
32
33   @Nullable
34   protected abstract PyExpression getCondition(PyStatement statement);
35
36   @Nullable
37   protected abstract PyStatementListContainer getStatementListContainer(PyStatement statement);
38
39   @Override
40   public TextRange surroundExpression(@NotNull Project project, @NotNull Editor editor, @NotNull PyExpression expression)
41     throws IncorrectOperationException {
42     TextRange currentCaretPosition = TextRange.from(editor.getCaretModel().getOffset(), 0);
43     PyStatement statement = PyElementGenerator.getInstance(project).
44       createFromText(LanguageLevel.getDefault(), PyStatement.class, getTextToGenerate());
45     final PyExpression condition = getCondition(statement);
46     if (condition == null) {
47       return currentCaretPosition;
48     }
49     condition.replace(expression);
50     statement = (PyStatement)CodeStyleManager.getInstance(project).reformat(statement);
51     statement = (PyStatement)expression.getParent().replace(statement);
52     PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(editor.getDocument());
53     PyStatementListContainer statementListContainer = getStatementListContainer(statement);
54     if (statementListContainer == null) {
55       return currentCaretPosition;
56     }
57     PyStatementList statementList = statementListContainer.getStatementList();
58     PyStatement[] statements = statementList.getStatements();
59     final TextRange range = statements[0].getTextRange();
60     editor.getDocument().deleteString(range.getStartOffset(), range.getEndOffset());
61     return TextRange.from(range.getStartOffset(), 0);
62   }
63
64   @Override
65   public boolean isApplicable(@NotNull PyExpression expr) {
66     return true;
67   }
68 }