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);
+  }
+}
+