From b7b902a21961aae050ea2970207878e741414204 Mon Sep 17 00:00:00 2001
From: yida_young <418621207@qq.com>
Date: Fri, 23 Oct 2020 18:27:41 +0800
Subject: [PATCH 1/2] Add code format for import keyword completion

---
 .../completion/PbCompletionContributor.java   | 25 ++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/idea/plugin/protoeditor/lang/completion/PbCompletionContributor.java b/core/src/main/java/idea/plugin/protoeditor/lang/completion/PbCompletionContributor.java
index e9d6219..1fb9021 100644
--- a/core/src/main/java/idea/plugin/protoeditor/lang/completion/PbCompletionContributor.java
+++ b/core/src/main/java/idea/plugin/protoeditor/lang/completion/PbCompletionContributor.java
@@ -18,6 +18,8 @@
 import com.intellij.codeInsight.completion.*;
 import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.codeInsight.lookup.LookupElementBuilder;
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiErrorElement;
 import com.intellij.psi.util.PsiTreeUtil;
@@ -232,6 +234,27 @@ protected void addCompletions(
   }
 
   private static LookupElement lookupElementWithSpace(String keyword) {
-    return LookupElementBuilder.create(keyword).withInsertHandler(AddSpaceInsertHandler.INSTANCE);
+    boolean needQuotas = "import".equalsIgnoreCase(keyword);
+    return LookupElementBuilder.create(keyword).withInsertHandler(
+            needQuotas ? new PbKeywordInsertHandler() : AddSpaceInsertHandler.INSTANCE);
+  }
+
+  private static class PbKeywordInsertHandler extends BasicInsertHandler<LookupElement> {
+
+    @Override
+    public void handleInsert(@NotNull InsertionContext context, @NotNull LookupElement item) {
+      super.handleInsert(context, item);
+      Editor editor = context.getEditor();
+
+      Document document = editor.getDocument();
+      if (!document.getText().substring(context.getTailOffset()).startsWith(";")) {
+        document.insertString(context.getTailOffset(), ";");
+      }
+
+      int offset = editor.getCaretModel().getOffset();
+      document.insertString(offset, " \"");
+      document.insertString(offset + 2, "\"");
+      editor.getCaretModel().moveToOffset(offset + 2);
+    }
   }
 }

From fbf20aa27bb6713053cb9ed494cacc18338b81b8 Mon Sep 17 00:00:00 2001
From: yida_young <418621207@qq.com>
Date: Fri, 23 Oct 2020 18:46:33 +0800
Subject: [PATCH 2/2] Fix testcase

---
 .../protoeditor/lang/completion/PbCompleteKeywordsTest.java     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/itest/java/idea/plugin/protoeditor/lang/completion/PbCompleteKeywordsTest.java b/core/src/itest/java/idea/plugin/protoeditor/lang/completion/PbCompleteKeywordsTest.java
index 8670fae..4a25555 100644
--- a/core/src/itest/java/idea/plugin/protoeditor/lang/completion/PbCompleteKeywordsTest.java
+++ b/core/src/itest/java/idea/plugin/protoeditor/lang/completion/PbCompleteKeywordsTest.java
@@ -61,7 +61,7 @@ public void testOptionAtTopLevel() {
   public void testImportAtTopLevel() {
     setInput("imp" + CARET_TAG);
     assertTrue(completeWithUniqueChoice());
-    assertResult("import ");
+    assertResult("import \"\";");
   }
 
   public void testMessageInMessageAutoInsertSpace() {