python/src/com/jetbrains/python/inspections/quickfix/MoveFromFutureImportQuickFix.java

   1 /*
   2  * Copyright 2000-2014 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.inspections.quickfix;
  17
  18 import com.intellij.codeInspection.LocalQuickFix;
  19 import com.intellij.codeInspection.ProblemDescriptor;
  20 import com.intellij.openapi.project.Project;
  21 import com.intellij.psi.PsiElement;
  22 import com.intellij.psi.PsiFile;
  23 import com.jetbrains.python.PyBundle;
  24 import com.jetbrains.python.codeInsight.imports.AddImportHelper;
  25 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
  26 import com.jetbrains.python.formatter.PyBlock;
  27 import com.jetbrains.python.psi.PyFile;
  28 import com.jetbrains.python.psi.PyStringLiteralExpression;
  29 import org.jetbrains.annotations.NotNull;
  30
  31 /**
  32  * @author Alexey.Ivanov
  33  */
  34 public class MoveFromFutureImportQuickFix implements LocalQuickFix {
  35   @NotNull
  36   public String getName() {
  37     return PyBundle.message("QFIX.move.from.future.import");
  38   }
  39
  40   @NotNull
  41   public String getFamilyName() {
  42     return getName();
  43   }
  44
  45   public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
  46     PsiElement problemElement = descriptor.getPsiElement();
  47     PsiFile psiFile = problemElement.getContainingFile();
  48     if (psiFile instanceof PyFile) {
  49       problemElement.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true);
  50       PyFile file = (PyFile)psiFile;
  51       PyStringLiteralExpression docString = DocStringUtil.findDocStringExpression(file);
  52       if (docString != null) {
  53         file.addAfter(problemElement, docString.getParent() /* PyExpressionStatement */);
  54       }
  55       else {
  56         file.addBefore(problemElement, file.getStatements().get(0));
  57       }
  58       problemElement.delete();
  59     }
  60   }
  61 }