feat: support name suggestions for other languages (closes #615)

carlrobertoh · Sep 10, 2024 · 9bbebc1 · 9bbebc1
1 parent 4ebe352
commit 9bbebc1
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 7 deletions.
diff --git a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java
@@ -24,4 +24,6 @@ public class CodeGPTKeys {
       Key.create("codegpt.editorCompletionLines");
   public static final Key<Boolean> IS_FETCHING_COMPLETION =
       Key.create("codegpt.isFetchingCompletion");
+  public static final Key<Boolean> IS_PROMPT_TEXT_FIELD_DOCUMENT =
+      Key.create("codegpt.isPromptTextFieldDocument");
 }
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/MethodNameLookupListener.java b/src/main/java/ee/carlrobert/codegpt/completions/MethodNameLookupListener.java
@@ -1,5 +1,7 @@
 package ee.carlrobert.codegpt.completions;
 
+import static ee.carlrobert.codegpt.CodeGPTKeys.IS_PROMPT_TEXT_FIELD_DOCUMENT;
+
 import com.intellij.codeInsight.completion.PrefixMatcher;
 import com.intellij.codeInsight.lookup.Lookup;
 import com.intellij.codeInsight.lookup.LookupElementBuilder;
@@ -8,19 +10,27 @@
 import com.intellij.openapi.application.Application;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.PsiUtilCore;
 import ee.carlrobert.codegpt.Icons;
 import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings;
+import java.util.List;
 import java.util.Optional;
 import org.jetbrains.annotations.Nullable;
 
 public class MethodNameLookupListener implements LookupManagerListener {
 
   @Override
   public void activeLookupChanged(@Nullable Lookup oldLookup, @Nullable Lookup newLookup) {
-    if (!ConfigurationSettings.getState().getMethodNameGenerationEnabled()
-        || !CompletionRequestService.isRequestAllowed()
-        || !(newLookup instanceof LookupImpl lookup)) {
+    if (!(newLookup instanceof LookupImpl lookup)) {
+      return;
+    }
+
+    Boolean isPromptTextFieldDocument = IS_PROMPT_TEXT_FIELD_DOCUMENT.get(
+        lookup.getEditor().getDocument());
+    if ((isPromptTextFieldDocument != null && isPromptTextFieldDocument)
+        || !ConfigurationSettings.getState().getMethodNameGenerationEnabled()
+        || !CompletionRequestService.isRequestAllowed()) {
       return;
     }
 
@@ -30,7 +40,7 @@ public void activeLookupChanged(@Nullable Lookup oldLookup, @Nullable Lookup new
         .ifPresent(context ->
             application.runReadAction(() -> {
               var type = PsiUtilCore.getElementType(context);
-              if ("METHOD".equals(type.toString())) {
+              if (PSIMethodMapping.contains(type)) {
                 var selection = context.getText();
                 application.executeOnPooledThread(() ->
                     addCompletionLookupValues(lookup, application, selection));
@@ -45,13 +55,38 @@ private void addCompletionLookupValues(
     CompletionRequestService.getInstance().getLookupCompletion(prompt)
         .ifPresent(response -> {
           for (var value : response.split(",")) {
-            application.runReadAction(() -> {
+            application.invokeLater(() -> application.runReadAction(() -> {
               lookup.addItem(
                   LookupElementBuilder.create(value.trim()).withIcon(Icons.Sparkle),
                   PrefixMatcher.ALWAYS_TRUE);
-              application.invokeLater(() -> lookup.refreshUi(true, true));
-            });
+              lookup.refreshUi(true, true);
+            }));
           }
         });
   }
+
+  private enum PSIMethodMapping {
+    GO(List.of("FILE", "METHOD_DECLARATION|FUNCTION_DECLARATION")),
+    JAVA(List.of("java.FILE", "METHOD")),
+    PY(List.of("FILE", "Py:FUNCTION_DECLARATION")),
+    CS(List.of("FILE", "DUMMY_TYPE_DECLARATION", "DUMMY_BLOCK")),
+    PHP(List.of("FILE", "Class method|function|Function")),
+    KOTLIN(List.of("FUN")),
+    DEFAULT(List.of("FILE", "METHOD_DECLARATION"));
+
+    private final List<String> types;
+
+    PSIMethodMapping(List<String> codes) {
+      this.types = codes;
+    }
+
+    public static boolean contains(IElementType type) {
+      for (var value : PSIMethodMapping.values()) {
+        if (value.types.contains(type.toString())) {
+          return true;
+        }
+      }
+      return false;
+    }
+  }
 }
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/PromptTextField.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/PromptTextField.kt
@@ -15,6 +15,7 @@ import com.intellij.ui.EditorTextField
 import com.intellij.ui.JBColor
 import com.intellij.util.ui.JBUI
 import ee.carlrobert.codegpt.CodeGPTBundle
+import ee.carlrobert.codegpt.CodeGPTKeys.IS_PROMPT_TEXT_FIELD_DOCUMENT
 import ee.carlrobert.codegpt.ui.textarea.suggestion.SuggestionsPopupManager
 import ee.carlrobert.codegpt.ui.textarea.suggestion.item.SuggestionActionItem
 import ee.carlrobert.codegpt.ui.textarea.suggestion.item.SuggestionItem
@@ -49,6 +50,7 @@ class PromptTextField(
         background = JBColor.background()
         minimumSize = Dimension(100, 40)
         document.addDocumentListener(getDocumentListener(onTextChanged))
+        IS_PROMPT_TEXT_FIELD_DOCUMENT.set(document, true)
         setPlaceholder(CodeGPTBundle.get("toolwindow.chat.textArea.emptyText"))
         IdeEventQueue.getInstance().addDispatcher(
             PromptTextFieldEventDispatcher(this, suggestionsPopupManager, appliedInlays) {