diff options
Diffstat (limited to 'python/src/com/jetbrains/python/PyAddImportFix.java')
-rw-r--r-- | python/src/com/jetbrains/python/PyAddImportFix.java | 55 |
1 files changed, 55 insertions, 0 deletions
diff --git a/python/src/com/jetbrains/python/PyAddImportFix.java b/python/src/com/jetbrains/python/PyAddImportFix.java new file mode 100644 index 000000000000..c13fed13fdcb --- /dev/null +++ b/python/src/com/jetbrains/python/PyAddImportFix.java @@ -0,0 +1,55 @@ +package com.jetbrains.python; + +import com.intellij.codeInspection.LocalQuickFix; +import com.intellij.codeInspection.ProblemDescriptor; +import com.intellij.openapi.project.Project; +import com.intellij.psi.PsiElement; +import com.jetbrains.python.codeInsight.imports.AddImportHelper; +import com.jetbrains.python.psi.LanguageLevel; +import com.jetbrains.python.psi.PyElementGenerator; +import com.jetbrains.python.psi.PyFile; +import com.jetbrains.python.psi.PyImportStatementBase; +import org.jetbrains.annotations.NotNull; + +/** + * Quick fix that adds import to file + * + * @author Ilya.Kazakevich + */ +public class PyAddImportFix implements LocalQuickFix { + @NotNull + private final String myImportToAdd; + @NotNull + private final PyFile myFile; + + /** + * @param importToAdd string representing what to add (i.e. "from foo import bar") + * @param file where to add + */ + public PyAddImportFix(@NotNull final String importToAdd, @NotNull final PyFile file) { + myImportToAdd = importToAdd; + myFile = file; + } + + @NotNull + @Override + public String getName() { + return PyBundle.message("QFIX.add.import", myImportToAdd); + } + + @NotNull + @Override + public String getFamilyName() { + return getName(); + } + + @Override + public void applyFix(@NotNull final Project project, @NotNull final ProblemDescriptor descriptor) { + final PyElementGenerator generator = PyElementGenerator.getInstance(project); + final PyImportStatementBase statement = + generator.createFromText(LanguageLevel.forElement(myFile), PyImportStatementBase.class, myImportToAdd); + final PsiElement recommendedPosition = AddImportHelper.getFileInsertPosition(myFile); + myFile.addAfter(statement, recommendedPosition); + } +} + |