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 }