Created a Goban class to handle all of the Go logic

pygo.py

@@ -42,8 +42,8 @@ def build_img_res():
 
     for d in ('u', 'd', 'l', 'r', 'm', 'dl', 'dr', 'ul', 'ur', 'h'):
         ret[d] = load_png('go_' + d + '.png')
-        ret[d + 'c'] = load_png('go_' + d + '.png')
+        ret[d + 'C'] = load_png('go_' + d + '.png')
-        ret[d + 'c'].blit(circle, (0,0))
+        ret[d + 'C'].blit(circle, (0,0))
 
     return ret
 
@@ -53,7 +53,8 @@ class GobanSquare:
 
     def __init__(self, pos):
         self.x, self.y = pos
-        self.state = 'empty'
+        self.state = -1
+        self.marked = False
 
         if (self.x, self.y) == (1,1):
             self.default_draw_code = 'ul'
@@ -77,34 +78,176 @@ class GobanSquare:
             self.default_draw_code = 'm'
 
 
+    def toggle_marked(self):
+        if self.marked:
+            self.marked = False
+        else:
+            self.marked = True
+
+
     def get_draw_code(self):
-        if self.state == 'empty':
+        ret = None
-            return self.default_draw_code
-        if self.state == 'marked':
-            return self.default_draw_code + 'c'
-        if self.state == 'white':
-            return 'w'
-        if self.state == 'black':
-            return 'b'
 
-        return None
+        if self.state == -1:
+            ret = self.default_draw_code
+        elif self.state == 0:
+            ret = 'b'
+        elif self.state == 1:
+            ret = 'w'
+
+        if self.marked:
+            if self.state >=0:
+                ret = self.default_draw_code + 'T'
+            else:
+                ret = self.default_draw_code + 'C'
+
+        return ret
 
 
 
-def draw_board(goban, size, img_res):
+class Goban:
-    board = pygame.Surface((size,size))
+    """Represents the go board. Handles stone placement, captures, etc"""
 
-    inc = size / 19;
+    def __init__(self):
-    i = 0
+        # Build the 361 board intersections
-    for row in goban:
+        self.board = []
-        j = 0
+        for i in range(19):
-        for square in row:
+            self.board.append([])
-            s = pygame.transform.scale(img_res[square.get_draw_code()], (inc, inc))
+            for j in range(19):
-            board.blit(s, (j*inc,i*inc))
+                self.board[i].append(GobanSquare((i+1, j+1)))
-            j += 1
+
-        i += 1
+        self.turn = 0
+
+
+    def place_stone(self, pos):
+        if not self._valid_move(pos):
+            return
+
+        x, y = pos
+        self.board[x][y].state = self.turn
+        self._capture(pos)
+        self.turn = (self.turn + 1) % 2
+
+
+    def toggle_marked(self, pos):
+        x,y = pos
+        if x < 0 or x > 18 or y < 0 or y > 18:
+            return
+        
+        self.board[x][y].toggle_marked()
+        
+
+    def _capture(self, pos):
+        x, y = pos
+        who = (self.turn + 1) % 2
+
+        if self._has_liberties((x, y+1), who):
+            self._delete_group((x, y+1))
+        if self._has_liberties((x, y-1), who):
+            self._delete_group((x, y-1))
+        if self._has_liberties((x+1, y), who):
+            self._delete_group((x+1, y))
+        if self._has_liberties((x-1, y), who):
+            self._delete_group((x-1, y))
+
+
+    def _valid_move(self, pos):
+        x, y = pos
+        liberties = self._has_liberties((x, y+1), self.turn, 'r')
+        liberties += self._has_liberties((x-1, y), self.turn, 'u')
+        liberties += self._has_liberties((x, y-1), self.turn, 'l')
+        liberties += self._has_liberties((x+1, y), self.turn, 'd')
+        return liberties
+
+
+    # Recursively find whether there are liberties for the group
+    # at pos. Positive numbers are not necessarily accurate
+    def _has_liberties(self, pos, who, direction = None):
+        x,y = pos
+        if x < 0 or x > 18 or y < 0 or y > 18:
+            return 0
+
+        square = self.board[x][y]
+        if square.state == -1:
+            return 1
+        elif square.state != who:
+            return 0
+        elif square.state == who:
+            if direction == None:
+                liberties = self._has_liberties((x, y+1), who, 'r')
+                liberties += self._has_liberties((x-1, y), who, 'u')
+                liberties += self._has_liberties((x+1, y), who, 'd')
+                liberties += self._has_liberties((x, y-1), who, 'l')
+            if direction == 'r':
+                liberties = self._has_liberties((x, y+1), who, 'r')
+                liberties += self._has_liberties((x-1, y), who, 'u')
+                liberties += self._has_liberties((x+1, y), who, 'd')
+            if direction == 'l':
+                liberties = self._has_liberties((x, y-1), who, 'l')
+                liberties += self._has_liberties((x-1, y), who, 'u')
+                liberties += self._has_liberties((x+1, y), who, 'd')
+            if direction == 'u':
+                liberties = self._has_liberties((x, y+1), who, 'r')
+                liberties += self._has_liberties((x-1, y), who, 'u')
+                liberties += self._has_liberties((x, y-1), who, 'l')
+            if direction == 'd':
+                liberties = self._has_liberties((x, y+1), who, 'r')
+                liberties += self._has_liberties((x, y-1), who, 'l')
+                liberties += self._has_liberties((x+1, y), who, 'd')
+
+            return liberties
+
+
+    # We don't need to worry about crossing ourselves with the
+    # recursion here, because we've already deleted the stone.
+    # It would be more efficient to avoid going backwards,
+    # but a lot more complicated (see _has_liberties)
+    def _delete_group(self, pos):
+        x,y = pos
+        if x < 0 or x > 18 or y < 0 or y > 18:
+            return
+
+        who = self.board[x][y].state
+
+        if who == -1:
+            return
+
+        self.board[x][y].state = -1
+
+        self._delete_group_r((x, y+1), who)
+        self._delete_group_r((x, y-1), who)
+        self._delete_group_r((x+1, y), who)
+        self._delete_group_r((x-1, y), who)
+
+
+    def _delete_group_r(self, pos, who):
+        x,y = pos
+        if x < 0 or x > 18 or y < 0 or y > 18 or self.board[x][y] != who:
+            return
+
+        self.board[x][y].state = -1
+
+        self._delete_group_r((x, y+1), who)
+        self._delete_group_r((x, y-1), who)
+        self._delete_group_r((x+1, y), who)
+        self._delete_group_r((x-1, y), who)
+
+
+    def draw_board(self, size, img_res):
+        ret = pygame.Surface((size,size))
+
+        inc = size / 19;
+        i = 0
+        for row in self.board:
+            j = 0
+            for square in row:
+                s = pygame.transform.scale(img_res[square.get_draw_code()], (inc, inc))
+                ret.blit(s, (j*inc,i*inc))
+                j += 1
+            i += 1
+
+        return ret.convert()
 
-    return board.convert()
 
 
 def main():
@@ -120,18 +263,13 @@ def main():
 
     # Build the dict of image objects
     img_res = build_img_res()
-
-    # Build the 361 board intersections
-    goban = []
-    for i in range(19):
-        goban.append([])
-        for j in range(19):
-            goban[i].append(GobanSquare((i+1, j+1)))
             
     board_size = 800
     board_inc = board_size / 19
 
-    board = draw_board(goban, board_size, img_res)
+    goban = Goban()
+
+    board = goban.draw_board(board_size, img_res)
     background.blit(board, (0,0))
     
     screen.blit(background, (0, 0))
@@ -143,17 +281,16 @@ def main():
         if event.type == QUIT:
             return
         if event.type == MOUSEBUTTONDOWN:
-            if event.button == 1:
-                state = 'black'
-            if event.button == 2:
-                state = 'marked'
-            if event.button == 3:
-                state = 'white'
-
             x, y = event.pos
-            goban[y / board_inc][x / board_inc].state = state
+            row = y / board_inc
-        
+            col = x / board_inc
-        board = draw_board(goban, board_size, img_res)
+
+            if event.button == 1:
+                goban.place_stone((row, col))
+            if event.button == 3:
+                goban.toggle_marked((row, col))
+                
+        board = goban.draw_board(board_size, img_res)
         background.blit(board, (0,0))
     
         screen.blit(background, (0, 0))