Code style fixes

Note: This will break coloring at the moment since e621 uses more tag types than Danbooru. I'll try to implement custom color support soon.

README.md

@@ -9,6 +9,8 @@ I created this script as a convenience tool since it reduces the need of switchi
 
 You can either download the files manually as described below, or use a pre-packaged version from [Releases](https://github.com/DominikDoom/a1111-sd-webui-tagcomplete/releases).
 
+(Note: e621 tags aren't added to the releases yet since coloring is broken for them at the moment).
+
 ### Disclaimer:
 This script is definitely not optimized, and it's not very intelligent. The tags are simply recommended based on their natural order in the CSV, which is their respective image count for the default Danbooru tag list. Also, at least for now, neither keyboard selection for tags nor completion for negative or img2img prompt textboxes is supported, and there's no way to turn the feature off from the ui, but I plan to get around to those features eventually.
 

tagAutocomplete.js

@@ -20,6 +20,9 @@ const autocompleteCSS_dark = `
     #autocompleteResultsList > li:hover {
         background-color: #1f2937;
     }
+    #autocompleteResultsList > li.selected {
+        background-color: #374151;
+    }
 `;
 const autocompleteCSS_light = `
     #autocompleteResults {
@@ -42,6 +45,9 @@ const autocompleteCSS_light = `
     #autocompleteResultsList > li:hover {
         background-color: #f5f6f8;
     }
+    #autocompleteResultsList > li.selected {
+        background-color: #e5e7eb;
+    }
 `;
 
 var acConfig = null;
@@ -112,6 +118,13 @@ const debounce = (func, wait = 300) => {
 
 // Difference function to fix duplicates not being seen as changes in normal filter
 function difference(a, b) {
+    if (a.length == 0) {
+        return b;
+    }
+    if (b.length == 0) {
+        return a;
+    }
+
     return [...b.reduce( (acc, v) => acc.set(v, (acc.get(v) || 0) - 1),
             a.reduce( (acc, v) => acc.set(v, (acc.get(v) || 0) + 1), new Map() ) 
     )].reduce( (acc, [v, count]) => acc.concat(Array(Math.abs(count)).fill(v)), [] );
@@ -129,14 +142,21 @@ function createResultsDiv() {
     return resultsDiv;
 }
 
+// The selected tag index. Needs to be up here so hide can access it.
+var selectedTag = null;
+
 // Show or hide the results div
+var isVisible = false;
 function showResults() {
     let resultsDiv = gradioApp().querySelector('#autocompleteResults');
     resultsDiv.style.display = "block";
+    isVisible = true;
 }
 function hideResults() {
     let resultsDiv = gradioApp().querySelector('#autocompleteResults');
     resultsDiv.style.display = "none";
+    isVisible = false;
+    selectedTag = null;
 }
 
 // On click, insert the tag into the prompt textbox with respect to the cursor position
@@ -144,11 +164,17 @@ function insertTextAtCursor(text, tagword) {
     let promptTextbox = gradioApp().querySelector('#txt2img_prompt > label > textarea');
     let cursorPos = promptTextbox.selectionStart;
     let sanitizedText = acConfig.replaceUnderscores ? text.replaceAll("_", " ") : text;
-    let optionalComma = (promptTextbox.value[cursorPos] == ",") ? "" : ", ";
+    
+    var prompt = promptTextbox.value;
+    let optionalComma = (prompt[cursorPos] === "," || prompt[cursorPos + tagword.length] === ",") ? "" : ", ";
 
     // Edit prompt text
-    var prompt = promptTextbox.value;
-    promptTextbox.value = prompt.substring(0, cursorPos - tagword.length) + sanitizedText + optionalComma + prompt.substring(cursorPos);
+    let direction = prompt.substring(cursorPos, cursorPos + tagword.length) === tagword ? 1 : -1;
+    if (direction === 1) {
+        promptTextbox.value = prompt.substring(0, cursorPos) + sanitizedText + optionalComma + prompt.substring(cursorPos + tagword.length)
+    } else {  
+        promptTextbox.value = prompt.substring(0, cursorPos - tagword.length) + sanitizedText + optionalComma + prompt.substring(cursorPos)
+    }
     prompt = promptTextbox.value;
 
     // Update cursor position to after the inserted text
@@ -163,30 +189,52 @@ function insertTextAtCursor(text, tagword) {
     previousTags = tags;
 }
 
-const colors_dark = ["lightblue", "indianred", "unused", "violet", "lightgreen", "orange"];
-const colors_light = ["dodgerblue", "firebrick", "unused", "darkorchid", "darkgreen", "darkorange" ]
 function addResultsToList(results, tagword) {
     let resultsList = gradioApp().querySelector('#autocompleteResultsList');
     resultsList.innerHTML = "";
 
-    let colors = gradioApp().querySelector('.dark') ? colors_dark : colors_light;
+    // Find right colors from config
+    let tagFileName = acConfig.tagFile.split(".")[0];
+    let tagColors = acConfig.colors;
+    //let colorIndex = Object.keys(tagColors).findIndex(key => key === tagFileName);
+    //let colorValues = Object.values(tagColors)[colorIndex];
+
+    let mode = gradioApp().querySelector('.dark') ? 0 : 1;
 
     for (let i = 0; i < results.length; i++) {
         let result = results[i];
         let li = document.createElement("li");
         li.innerHTML = result[0];
-        li.style = `color: ${colors[result[1]]};`;
+
+        // Set the color of the tag
+        let tagType = result[1];
+        li.style = `color: ${tagColors[tagFileName][tagType][mode]};`;
+        // Add listener
         li.addEventListener("click", function() { insertTextAtCursor(result[0], tagword); });
+        // Add element to list
         resultsList.appendChild(li);
     }
 }
 
+function updateSelectionStyle(num) {
+    let resultsList = gradioApp().querySelector('#autocompleteResultsList');
+    let items = resultsList.getElementsByTagName('li');
+
+    for (let i = 0; i < items.length; i++) {
+        items[i].classList.remove('selected');
+    }
+
+    items[num].classList.add('selected');
+}
+
 allTags = [];
 previousTags = [];
-
+results = [];
+tagword = "";
+resultCount = 0;
 function autocomplete(prompt) {
     // Guard for empty prompt
-    if (prompt.length == 0) {
+    if (prompt.length === 0) {
         hideResults();
         return;
     }
@@ -197,23 +245,24 @@ function autocomplete(prompt) {
     previousTags = tags;
 
     // Guard for no difference / only whitespace remaining
-    if (diff == undefined || diff.length == 0) {
+    if (diff === undefined || diff.length === 0) {
         hideResults();
         return;
     }
 
-    let tagword = diff[0]
+    tagword = diff[0]
 
     // Guard for empty tagword
-    if (tagword == undefined || tagword.length == 0) {
+    if (tagword === undefined || tagword.length === 0) {
         hideResults();
         return;
     }
     
-    let results = allTags.filter(x => x[0].includes(tagword)).slice(0, acConfig.maxResults);
+    results = allTags.filter(x => x[0].includes(tagword)).slice(0, acConfig.maxResults);
+    resultCount = results.length;
 
     // Guard for empty results
-    if (results.length == 0) {
+    if (resultCount === 0) {
         hideResults();
         return;
     }
@@ -222,14 +271,59 @@ function autocomplete(prompt) {
     addResultsToList(results, tagword);
 }
 
+function navigateInList(event) {
+    validKeys = ["ArrowUp", "ArrowDown", "ArrowLeft", "ArrowRight", "Enter", "Escape"];
+
+    if (!validKeys.includes(event.key)) return;
+    if (!isVisible) return
+
+    switch (event.key) {
+        case "ArrowUp":
+            if (selectedTag === null) {
+                selectedTag = resultCount - 1;
+            } else {
+                selectedTag = (selectedTag - 1 + resultCount) % resultCount;
+            }
+            break;
+        case "ArrowDown":
+            if (selectedTag === null) {
+                selectedTag = 0;
+            } else {
+                selectedTag = (selectedTag + 1) % resultCount;
+            }
+            break;
+        case "ArrowLeft":
+            selectedTag = 0;
+            break;
+        case "ArrowRight":
+            selectedTag = resultCount - 1;
+            break;
+        case "Enter":
+            if (selectedTag !== null) {
+                insertTextAtCursor(results[selectedTag][0], tagword);
+            }
+            break;
+        case "Escape":
+            hideResults();
+            break;
+    }
+    // Update highlighting
+    if (selectedTag !== null)
+        updateSelectionStyle(selectedTag);
+
+    // Prevent default behavior
+    event.preventDefault();
+    event.stopPropagation();
+}
+
 onUiUpdate(function(){
     // One-time CSV setup
-    if (acConfig == null) acConfig = JSON.parse(readFile("file/tags/config.json"));
-    if (allTags.length == 0) allTags = loadCSV();
+    if (acConfig === null) acConfig = JSON.parse(readFile("file/tags/config.json"));
+    if (allTags.length === 0) allTags = loadCSV();
 
 	let promptTextbox = gradioApp().querySelector('#txt2img_prompt > label > textarea');
 	
-    if (promptTextbox == null) return;
+    if (promptTextbox === null) return;
     if (gradioApp().querySelector('#autocompleteResults') != null) return;
 
     // Only add listeners once
@@ -244,6 +338,10 @@ onUiUpdate(function(){
         promptTextbox.addEventListener('input', debounce(() => autocomplete(promptTextbox.value), 100));
         // Add focusout event listener
         promptTextbox.addEventListener('focusout', debounce(() => hideResults(), 400));
+        // Add up and down arrow event listener
+        promptTextbox.addEventListener('keydown', (e) => navigateInList(e));
+
+    
 
         // Add class so we know we've already added the listeners
         promptTextbox.classList.add('autocomplete');

tags/config.json

@@ -1,5 +1,25 @@
 {
 	"tagFile": "danbooru.csv",
 	"maxResults": 5,
-	"replaceUnderscores": true
-}
+	"replaceUnderscores": true,
+	"colors": {
+		"danbooru": {
+			"0": ["lightblue", "dodgerblue"],
+			"1": ["indianred", "firebrick"],
+			"3": ["violet", "darkorchid"],
+			"4": ["lightgreen", "darkgreen"],
+			"5": ["orange", "darkorange"]
+		},
+		"e621": {
+			"-1": ["red", "maroon"],
+			"0": ["lightblue", "dodgerblue"],
+			"1": ["gold", "goldenrod"],
+			"3": ["violet", "darkorchid"],
+			"4": ["lightgreen", "darkgreen"],
+			"5": ["tomato", "darksalmon"],
+			"6": ["red", "maroon"],
+			"7": ["whitesmoke", "black"],
+			"8": ["seagreen", "darkseagreen"]
+		}
+	}
+}