Improves suggestion text sanitization logic

aws · Oct 3, 2024 · 380712f · 380712f
1 parent c814d7e
commit 380712f
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 2 deletions.
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/util/QInvocationSession.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/util/QInvocationSession.java
@@ -28,6 +28,7 @@
 
 import static software.aws.toolkits.eclipse.amazonq.util.QConstants.Q_INLINE_HINT_TEXT_STYLE;
 import static software.aws.toolkits.eclipse.amazonq.util.QEclipseEditorUtils.getActiveTextViewer;
+import static software.aws.toolkits.eclipse.amazonq.util.SuggestionTextUtil.replaceSpacesWithTabs;
 
 public final class QInvocationSession extends QResource {
 
@@ -160,8 +161,8 @@ public void invoke() {
                             .inlineCompletionWithReferences(params)
                             .thenApply(result -> result.getItems().parallelStream().map(item -> {
                                 if (isTabOnly) {
-                                    String origText = item.getInsertText();
-                                    String sanitizedText = origText.replace("\n" + " ".repeat(tabSize), "\n\t");
+                                    String sanitizedText = replaceSpacesWithTabs(item.getInsertText(), tabSize);
+                                    System.out.println("Sanitized text: " + sanitizedText.replace("\n", "\\n").replace("\t", "\\t"));
                                     item.setInsertText(sanitizedText);
                                 }
                                 return item;

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/util/SuggestionTextUtil.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/util/SuggestionTextUtil.java
@@ -0,0 +1,55 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+package software.aws.toolkits.eclipse.amazonq.util;
+
+public final class SuggestionTextUtil {
+
+    private SuggestionTextUtil() {
+    }
+
+    public static String replaceSpacesWithTabs(String input, int tabSize) {
+        StringBuilder result = new StringBuilder();
+        String[] lines = input.split("\\r?\\n");
+
+        for (int i = 0; i < lines.length; i++) {
+            String line = lines[i];
+            int numSpaces = 0;
+            StringBuilder newLine = new StringBuilder();
+
+            for (int j = 0; j < line.length(); j++) {
+                char c = line.charAt(j);
+                if (c == ' ') {
+                    numSpaces++;
+                } else {
+                    newLine.append(getTabsForSpaces(numSpaces, tabSize));
+                    newLine.append(c);
+                    numSpaces = 0;
+                }
+            }
+
+            if (i < lines.length - 1) {
+                newLine.append("\n");
+            }
+
+            result.append(newLine);
+        }
+
+        return result.toString();
+    }
+
+    private static String getTabsForSpaces(int numSpaces, int tabSize) {
+        int numTabs = numSpaces / tabSize;
+        StringBuilder tabs = new StringBuilder();
+
+        for (int i = 0; i < numTabs; i++) {
+            tabs.append("\t");
+        }
+
+        int remainingSpaces = numSpaces % tabSize;
+        for (int i = 0; i < remainingSpaces; i++) {
+            tabs.append(" ");
+        }
+
+        return tabs.toString();
+    }
+}