Add factories_by_player, handle non-movement actions in expected_position.

1 files changed, 14 insertions, 4 deletions

diff --git a/robots/state.py b/robots/state.py
index 8a47139..453cb7b 100644
--- a/robots/state.py
+++ b/robots/state.py
@@ -86,6 +86,14 @@ class GameState:
         }
 
     @property
+    def factories_by_player(self):
+        result = defaultdict(list)
+        for p in self.factories:
+            player = self.allegiances.get(p)
+            result[player].append(p)
+        return result
+
+    @property
     def n_alive_players(self):
         """How many players are still alive."""
         return sum(
@@ -121,12 +129,14 @@ class GameState:
         Of course, combat (and possibly other factors in the future) may
         influence whether it actually ends up at the returned position.
         """
+        if action not in DIRECTIONS:
+            return (x, y)
+
         dx, dy = DIRECTIONS[action]
         nx = (x + dx) % self.width
         ny = (y + dy) % self.height
 
-        if self.cities[ny][nx] in City.traversable:
-            x = nx
-            y = ny
+        if self.cities[ny][nx] not in City.traversable:
+            return (x, y)
 
-        return (x, y)
+        return (nx, ny)