From d673a521d86faedbd86aa63a3912c1191bb486d1 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Tue, 8 Oct 2024 20:12:23 -0400
Subject: [PATCH] Add always snap to grid setting (#1177)

* Always snap to grid

* Ban pysssss.SnapToGrid

* nit
---
 src/extensions/core/snapToGrid.ts | 27 ++++++++++++++++++++-------
 src/stores/extensionStore.ts      |  5 +++++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/extensions/core/snapToGrid.ts b/src/extensions/core/snapToGrid.ts
index a4664201a..9a5e49954 100644
--- a/src/extensions/core/snapToGrid.ts
+++ b/src/extensions/core/snapToGrid.ts
@@ -1,3 +1,4 @@
+import type { SettingParams } from '@/types/settingTypes'
 import { app } from '../../scripts/app'
 import {
   LGraphCanvas,
@@ -37,13 +38,25 @@ app.registerExtension({
         LiteGraph.CANVAS_GRID_SIZE = +value || 10
       }
     })
+    // Keep the 'pysssss.SnapToGrid' setting id so we don't need to migrate setting values.
+    // Using a new setting id can cause existing users to lose their existing settings.
+    const alwaysSnapToGrid = app.ui.settings.addSetting({
+      id: 'pysssss.SnapToGrid',
+      category: ['Comfy', 'Graph', 'AlwaysSnapToGrid'],
+      name: 'Always snap to grid',
+      type: 'boolean',
+      defaultValue: false,
+      versionAdded: '1.3.13'
+    } as SettingParams)
+
+    const shouldSnapToGrid = () => app.shiftDown || alwaysSnapToGrid.value
 
     // After moving a node, if the shift key is down align it to grid
     const onNodeMoved = app.canvas.onNodeMoved
     app.canvas.onNodeMoved = function (node) {
       const r = onNodeMoved?.apply(this, arguments)
 
-      if (app.shiftDown) {
+      if (shouldSnapToGrid()) {
         // Ensure all selected nodes are realigned
         for (const id in this.selected_nodes) {
           this.selected_nodes[id].alignToGrid()
@@ -58,7 +71,7 @@ app.registerExtension({
     app.graph.onNodeAdded = function (node) {
       const onResize = node.onResize
       node.onResize = function () {
-        if (app.shiftDown) {
+        if (shouldSnapToGrid()) {
           roundVectorToGrid(node.size)
         }
         return onResize?.apply(this, arguments)
@@ -70,7 +83,7 @@ app.registerExtension({
     const origDrawNode = LGraphCanvas.prototype.drawNode
     LGraphCanvas.prototype.drawNode = function (node, ctx) {
       if (
-        app.shiftDown &&
+        shouldSnapToGrid() &&
         this.node_dragged &&
         node.id in this.selected_nodes
       ) {
@@ -132,8 +145,8 @@ app.registerExtension({
       // to snap on a mouse-up which we can determine by checking if `app.canvas.last_mouse_dragging`
       // has been set to `false`. Essentially, this check here is the equivalent to calling an
       // `LGraphGroup.prototype.onNodeMoved` if it had existed.
-      if (app.canvas.last_mouse_dragging === false && app.shiftDown) {
-        // After moving a group (while app.shiftDown), snap all the child nodes and, finally,
+      if (app.canvas.last_mouse_dragging === false && shouldSnapToGrid()) {
+        // After moving a group (while shouldSnapToGrid()), snap all the child nodes and, finally,
         //  align the group itself.
         this.recomputeInsideNodes()
         for (const node of this.nodes) {
@@ -151,7 +164,7 @@ app.registerExtension({
      */
     const drawGroups = LGraphCanvas.prototype.drawGroups
     LGraphCanvas.prototype.drawGroups = function (canvas, ctx) {
-      if (this.selected_group && app.shiftDown) {
+      if (this.selected_group && shouldSnapToGrid()) {
         if (this.selected_group_resizing) {
           roundVectorToGrid(this.selected_group.size)
         } else if (selectedAndMovingGroup) {
@@ -176,7 +189,7 @@ app.registerExtension({
     const onGroupAdd = LGraphCanvas.onGroupAdd
     LGraphCanvas.onGroupAdd = function () {
       const v = onGroupAdd.apply(app.canvas, arguments)
-      if (app.shiftDown) {
+      if (shouldSnapToGrid()) {
         const lastGroup = app.graph.groups[app.graph.groups.length - 1]
         if (lastGroup) {
           roundVectorToGrid(lastGroup.pos)
diff --git a/src/stores/extensionStore.ts b/src/stores/extensionStore.ts
index e2d71382a..440c3147d 100644
--- a/src/stores/extensionStore.ts
+++ b/src/stores/extensionStore.ts
@@ -62,6 +62,11 @@ export const useExtensionStore = defineStore('extension', () => {
     // pysssss.Locking is replaced by pin/unpin in ComfyUI core.
     // https://github.com/Comfy-Org/litegraph.js/pull/117
     disabledExtensionNames.value.add('pysssss.Locking')
+    // pysssss.SnapToGrid is replaced by Comfy.Graph.AlwaysSnapToGrid in ComfyUI core.
+    // pysssss.SnapToGrid tries to write global app.shiftDown state, which is no longer
+    // allowed since v1.3.12.
+    // https://github.com/Comfy-Org/ComfyUI_frontend/issues/1176
+    disabledExtensionNames.value.add('pysssss.SnapToGrid')
   }
 
   // Some core extensions are registered before the store is initialized, e.g.