PY-17564 Keep blank lines between imports inserted by user
[idea/community.git] / 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 }