fix up map generation functions to use constants

1 files changed, 10 insertions, 7 deletions

diff --git a/robots/utils.py b/robots/utils.py
index dff2a3e..e269a02 100644
--- a/robots/utils.py
+++ b/robots/utils.py
@@ -1,6 +1,8 @@
 from random import sample
 import time
 
+from robots.constants import City
+
 def ceil_div(a, b):
     """Divide a by b, rounding towards infinity."""
     return -(-a // b)
@@ -8,15 +10,16 @@ def ceil_div(a, b):
 def ilen(items):
     return sum(1 for _ in items)
 
-def add_spawns(map_, n_spawns):
+def add_spawns(map_, n_spawns, city=None):
     available = []
     for x, y, cell in iter_board(map_):
-        if cell != '*':
+        if cell != City.GHOST:
             available.append((x, y))
 
     spawns = sample(available, n_spawns)
     for i, (x, y) in enumerate(spawns):
-        map_[y][x] = str(i)
+        map_[y][x] = city or str(i)
+
 
 def empty_map(width, height, n_spawns):
     board = [['.'] * width for y in range(height)]
@@ -30,11 +33,11 @@ def border_map(width, height, n_spawns):
         for y in range(height + 2)
     ]
     for x in range(width + 2):
-        board[0][x] = '*'
-        board[-1][x] = '*'
+        board[0][x] = City.GHOST
+        board[-1][x] = City.GHOST
     for y in range(height + 2):
-        board[y][0] = '*'
-        board[y][-1] = '*'
+        board[y][0] = City.GHOST
+        board[y][-1] = City.GHOST
 
     add_spawns(board, n_spawns)
     return board