Added a Diagonal constraint.

[pseudoku.git] / pseudoku / grid / cellgroup.py

diff --git a/pseudoku/grid/cellgroup.py b/pseudoku/grid/cellgroup.py
index 374bd20..841a490 100644 (file)
--- a/pseudoku/grid/cellgroup.py
+++ b/pseudoku/grid/cellgroup.py
@@ -84,3 +84,28 @@ class Column(CellConstraint):
         self._cells = []
         for row in xrange(self.grid._size):
             self._cells.append(ref(self.grid.cell(row, self._pos)))
+
+class Diagonal(CellConstraint):
+    def __init__(self, grid, direction='down', offset=0):
+        self._grid = ref(grid)
+        self._direction = direction
+        self._offset = offset
+
+        if direction == 'down':
+            coords = {'row': 0, 'column': 0}
+            increment = 1
+        elif direction == 'up':
+            coords = {'row': 0, 'column': grid.size - 1}
+            increment = -1
+        else:
+            raise Exception # XXX
+
+        # XXX
+        if offset != 0:
+            raise NotImplementedError
+
+        self._cells = []
+        for i in xrange(self.grid.size):
+            self._cells.append(ref(self.grid.cell(coords['row'], coords['column'])))
+            coords['row'] += 1
+            coords['column'] += increment