From 9da9b39d8b0f276fdff0f389a72d39ddcd8920ac Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Fri, 22 Jul 2022 16:07:07 -0400
Subject: [PATCH] Add on type formatting with snippet support

---
 package.json  |  8 +++++++-
 src/client.ts | 56 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/package.json b/package.json
index a3967ddf..e7639d24 100644
--- a/package.json
+++ b/package.json
@@ -79,6 +79,11 @@
               "type": "boolean",
               "default": true
             },
+            "onTypeFormatting": {
+              "description": "Enable on type formatting",
+              "type": "boolean",
+              "default": true
+            },
             "diagnostics": {
               "description": "Enable diagnostics, like RuboCop violations",
               "type": "boolean",
@@ -99,7 +104,8 @@
             "semanticHighlighting": true,
             "formatting": true,
             "diagnostics": true,
-            "codeActions": true
+            "codeActions": true,
+            "onTypeFormatting": true
           }
         }
       }
diff --git a/src/client.ts b/src/client.ts
index 9b8e6ec0..a9992c5c 100644
--- a/src/client.ts
+++ b/src/client.ts
@@ -60,6 +60,62 @@ export default class Client {
       initializationOptions: {
         enabledFeatures: this.listOfEnabledFeatures(),
       },
+      middleware: {
+        provideOnTypeFormattingEdits: async (
+          document,
+          position,
+          ch,
+          options,
+          token,
+          _next
+        ) => {
+          if (this.client) {
+            const response: vscode.TextEdit[] | null =
+              await this.client.sendRequest(
+                "textDocument/onTypeFormatting",
+                {
+                  textDocument: { uri: document.uri.toString() },
+                  position,
+                  ch,
+                  options,
+                },
+                token
+              );
+
+            if (!response) {
+              return null;
+            }
+
+            // Find the $0 anchor to move the cursor
+            const cursorPosition = response.find(
+              (edit) => edit.newText === "$0"
+            );
+
+            if (!cursorPosition) {
+              return response;
+            }
+
+            // Remove the edit including the $0 anchor
+            response.splice(response.indexOf(cursorPosition), 1);
+
+            const workspaceEdit = new vscode.WorkspaceEdit();
+            workspaceEdit.set(document.uri, response);
+            await vscode.workspace.applyEdit(workspaceEdit);
+
+            await vscode.window.activeTextEditor!.insertSnippet(
+              new vscode.SnippetString(cursorPosition.newText),
+              new vscode.Selection(
+                cursorPosition.range.start,
+                cursorPosition.range.end
+              )
+            );
+
+            return null;
+          }
+
+          return undefined;
+        },
+      },
     };
 
     this.context = context;