undoable Life component grid actions

1 files changed, 99 insertions, 33 deletions

diff --git a/src/components/life.py b/src/components/life.py
index d4a455d..7a610eb 100644
--- a/src/components/life.py
+++ b/src/components/life.py
@@ -1,4 +1,5 @@
 from PyQt5 import QtGui, QtCore, QtWidgets
+from PyQt5.QtWidgets import QUndoCommand
 from PIL import Image, ImageDraw, ImageEnhance, ImageChops, ImageFilter
 import os
 import math
@@ -58,22 +59,8 @@ class Component(Component):
             self.mergeUndo = True
 
     def shiftGrid(self, d):
-        def newGrid(Xchange, Ychange):
-            return {
-                (x + Xchange, y + Ychange)
-                for x, y in self.startingGrid
-            }
-
-        if d == 0:
-            newGrid = newGrid(0, -1)
-        elif d == 1:
-            newGrid = newGrid(0, 1)
-        elif d == 2:
-            newGrid = newGrid(-1, 0)
-        elif d == 3:
-            newGrid = newGrid(1, 0)
-        self.startingGrid = newGrid
-        self._sendUpdateSignal()
+        action = ShiftGrid(self, d)
+        self.parent.undoStack.push(action)
 
     def update(self):
         self.updateGridSize()
@@ -98,17 +85,14 @@ class Component(Component):
         enabled = (len(self.startingGrid) > 0)
         for widget in self.shiftButtons:
             widget.setEnabled(enabled)
-        super().update()
 
     def previewClickEvent(self, pos, size, button):
         pos = (
             math.ceil((pos[0] / size[0]) * self.gridWidth) - 1,
             math.ceil((pos[1] / size[1]) * self.gridHeight) - 1
         )
-        if button == 1:
-            self.startingGrid.add(pos)
-        elif button == 2:
-            self.startingGrid.discard(pos)
+        action = ClickGrid(self, pos, button)
+        self.parent.undoStack.push(action)
 
     def updateGridSize(self):
         w, h = self.core.resolutions[-1].split('x')
@@ -223,7 +207,7 @@ class Component(Component):
                         'up', 'down', 'left', 'right',
                     )
                 }
-                for cell in nearbyCoords(x, y):
+                for cell in self.nearbyCoords(x, y):
                     if cell not in grid:
                         continue
                     if cell[0] == x:
@@ -363,7 +347,7 @@ class Component(Component):
 
         def neighbours(x, y):
             return {
-                cell for cell in nearbyCoords(x, y)
+                cell for cell in self.nearbyCoords(x, y)
                 if cell in lastGrid
             }
 
@@ -374,7 +358,7 @@ class Component(Component):
                 newGrid.add((x, y))
         potentialNewCells = {
             coordTup for origin in lastGrid
-            for coordTup in list(nearbyCoords(*origin))
+            for coordTup in list(self.nearbyCoords(*origin))
         }
         for x, y in potentialNewCells:
             if (x, y) in newGrid:
@@ -397,13 +381,95 @@ class Component(Component):
                 widget.setEnabled(True)
         super().loadPreset(pr, *args)
 
+    def nearbyCoords(self, x, y):
+        yield x + 1, y + 1
+        yield x + 1, y - 1
+        yield x - 1, y + 1
+        yield x - 1, y - 1
+        yield x, y + 1
+        yield x, y - 1
+        yield x + 1, y
+        yield x - 1, y
+
+
+class ClickGrid(QUndoCommand):
+    def __init__(self, comp, pos, id_):
+        super().__init__(
+            "click %s component #%s" % (comp.name, comp.compPos))
+        self.comp = comp
+        self.pos = [pos]
+        self.id_ = id_
+
+    def id(self):
+        return self.id_
+
+    def mergeWith(self, other):
+        self.pos.extend(other.pos)
+        return True
+
+    def add(self):
+        for pos in self.pos[:]:
+            self.comp.startingGrid.add(pos)
+        self.comp.update(auto=True)
+
+    def remove(self):
+        for pos in self.pos[:]:
+            self.comp.startingGrid.discard(pos)
+        self.comp.update(auto=True)
+
+    def redo(self):
+        if self.id_ == 1:  # Left-click
+            self.add()
+        elif self.id_ == 2:  # Right-click
+            self.remove()
+
+    def undo(self):
+        if self.id_ == 1:  # Left-click
+            self.remove()
+        elif self.id_ == 2:  # Right-click
+            self.add()
+
+class ShiftGrid(QUndoCommand):
+    def __init__(self, comp, direction):
+        super().__init__(
+            "change %s component #%s" % (comp.name, comp.compPos))
+        self.comp = comp
+        self.direction = direction
+        self.distance = 1
+
+    def id(self):
+        return self.direction
+
+    def mergeWith(self, other):
+        self.distance += other.distance
+        return True
+
+    def newGrid(self, Xchange, Ychange):
+        return {
+            (x + Xchange, y + Ychange)
+            for x, y in self.comp.startingGrid
+        }
 
-def nearbyCoords(x, y):
-    yield x + 1, y + 1
-    yield x + 1, y - 1
-    yield x - 1, y + 1
-    yield x - 1, y - 1
-    yield x, y + 1
-    yield x, y - 1
-    yield x + 1, y
-    yield x - 1, y
+    def redo(self):
+        if self.direction == 0:
+            newGrid = self.newGrid(0, -self.distance)
+        elif self.direction == 1:
+            newGrid = self.newGrid(0, self.distance)
+        elif self.direction == 2:
+            newGrid = self.newGrid(-self.distance, 0)
+        elif self.direction == 3:
+            newGrid = self.newGrid(self.distance, 0)
+        self.comp.startingGrid = newGrid
+        self.comp._sendUpdateSignal()
+
+    def undo(self):
+        if self.direction == 0:
+            newGrid = self.newGrid(0, self.distance)
+        elif self.direction == 1:
+            newGrid = self.newGrid(0, -self.distance)
+        elif self.direction == 2:
+            newGrid = self.newGrid(self.distance, 0)
+        elif self.direction == 3:
+            newGrid = self.newGrid(-self.distance, 0)
+        self.comp.startingGrid = newGrid
+        self.comp._sendUpdateSignal()