From 362a8bfd2cdb853c0a184397140789cd4e3c35af Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Fri, 4 Aug 2023 20:39:44 +0100
Subject: [PATCH] Align to grid on resize

---
 js/snapToGrid.js | 49 +++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 40 insertions(+), 9 deletions(-)

diff --git a/js/snapToGrid.js b/js/snapToGrid.js
index ab417a5..0c905a8 100644
--- a/js/snapToGrid.js
+++ b/js/snapToGrid.js
@@ -1,10 +1,11 @@
 import { app } from "/scripts/app.js";
 
+let setting;
 const id = "pysssss.SnapToGrid";
 const ext = {
 	name: id,
-	async setup(app) {
-		app.ui.settings.addSetting({
+	init() {
+		setting = app.ui.settings.addSetting({
 			id,
 			name: "🐍 Always snap to grid",
 			defaultValue: false,
@@ -13,13 +14,43 @@ const ext = {
 				app.canvas.align_to_grid = value;
 			},
 		});
-		const origDrawNode = LGraphCanvas.prototype.drawNode;
-		LGraphCanvas.prototype.drawNode = function () {
-			const shift = app.shiftDown;
-			app.shiftDown = true; // Force it to draw the drop position
-			const r = origDrawNode.apply(this, arguments);
-			app.shiftDown = shift;
-            return r;
+
+		// We need to register our hooks after the core snap to grid extension runs
+		// Do this from the graph configure function so we still get onNodeAdded calls
+		const configure = LGraph.prototype.configure;
+		LGraph.prototype.configure = function () {
+			// Override drawNode to draw the drop position
+			const drawNode = LGraphCanvas.prototype.drawNode;
+			LGraphCanvas.prototype.drawNode = function () {
+				if (setting?.value) {
+					const shift = app.shiftDown;
+					app.shiftDown = true; 
+					const r = drawNode.apply(this, arguments);
+					app.shiftDown = shift;
+					return r;
+				}
+				return drawNode.apply(this, arguments);
+			};
+
+			// Override node added to add a resize handler to force grid alignment
+			const onNodeAdded = app.graph.onNodeAdded;
+			app.graph.onNodeAdded = function (node) {
+				const r = onNodeAdded.apply(this, arguments);
+				const onResize = node.onResize;
+				node.onResize = function () {
+					if (setting?.value) {
+						const shift = app.shiftDown;
+						app.shiftDown = true;
+						const r = onResize?.apply(this, arguments);
+						app.shiftDown = shift;
+						return r;
+					}
+					return onResize?.apply(this, arguments);
+				};
+				return r;
+			};
+
+			return configure.apply(this, arguments);
 		};
 	},
 };