Merge pull request #16 from profiq/class-docstrings

Implement better support for documenting classes
profiq · Jan 3, 2024 · bf957fd · bf957fd
2 parents 9d65349 + 007eb4d
commit bf957fd
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -4,7 +4,7 @@ plugins {
 }
 
 group = "com.profiq"
-version = "0.1.3"
+version = "0.1.4"
 
 repositories {
     mavenCentral()
@@ -19,7 +19,7 @@ dependencies {
 // Read more: https://plugins.jetbrains.com/docs/intellij/tools-gradle-intellij-plugin.html
 intellij {
     version.set("2023.3")
-    type.set("IC") // Target IDE Platform
+    type.set("PC") // Target IDE Platform
 
     plugins.set(listOf(/* Plugin Dependencies */))
 }

diff --git a/src/main/java/com/profiq/codexor/GenerateDocumentationAction.java b/src/main/java/com/profiq/codexor/GenerateDocumentationAction.java
@@ -56,7 +56,7 @@ public void actionPerformed(@NotNull AnActionEvent e) {
         PsiElement element = e.getData(CommonDataKeys.PSI_ELEMENT);
 
         if (element == null) {
-            showError("Please select a Python function");
+            showError("Please select a Python function, method or class");
             return;
         }
 
@@ -127,7 +127,10 @@ private String parseDocstring(String responseText) {
         Pattern pattern = Pattern.compile(":\n(\\s*\"\"\".*\"\"\")", Pattern.DOTALL);
         Matcher matcher = pattern.matcher(responseText);
         if (matcher.find()) {
-            return matcher.group(1);
+            String docstring = matcher.group(1);
+            String[] docstringSplit = docstring.split("\"\"\"");
+            docstring = docstringSplit[0] + "\"\"\"" + docstringSplit[1] + "\"\"\"";
+            return docstring;
         } else {
             return "";
         }
@@ -187,7 +190,6 @@ private void showEditor(AnActionEvent e, String docstring) {
         editor.getContentComponent().setPreferredSize(new Dimension(1000, 800));
         var confirmBtn = new JButton("Insert");
 
-
         Editor mainEditor = e.getData(CommonDataKeys.EDITOR);
         Document mainDocument = mainEditor.getDocument();
         LogicalPosition caretPosition = mainEditor.getCaretModel().getLogicalPosition();
@@ -225,6 +227,7 @@ private String applyIndentation(String docstring, String indentation) {
 
             for (String line : lines) {
                 unindented.append(indentation);
+
                 if (line.length() >= minIndent) {
                     unindented.append(line.substring(minIndent));
                 }
@@ -242,7 +245,7 @@ private String determineCorrectIndentation(String code) {
         Matcher matcher = pattern.matcher(code);
 
         if(matcher.find()) {
-            return matcher.group(1);
+            return matcher.group(1).replace("\n", "");
         }
 
         return "";

diff --git a/src/main/java/com/profiq/codexor/SettingsConfigurable.java b/src/main/java/com/profiq/codexor/SettingsConfigurable.java
@@ -16,7 +16,7 @@ public class SettingsConfigurable implements Configurable {
     public static String MODEL_SETTINGS_KEY = "codexor.model";
     public static String MODEL_DEFAULT = "gpt-3.5-turbo";
     public static String PROMPT_SETTINGS_KEY = "codexor.prompt";
-    public static String PROMPT_DEFAULT = "Generate high-quality docstring for the following Python function including function signature: ";
+    public static String PROMPT_DEFAULT = "Generate high-quality docstring for the following Python class or function.  Use the reStructuredText docstring format. Only return the docstring for the top-level element:";
 
     private final JTextField apiKeyField;
     private final JTextField modelField;