Add roach_at helper.

[koperkapel.git] / koperkapel / vehicles / base.py

diff --git a/koperkapel/vehicles/base.py b/koperkapel/vehicles/base.py
index 9aff447f1fdbb0134635059671773734a9a6bd9a..c2a24f64ecbfb6cb5d9f27cc03ae11da0a3cf141 100644 (file)
--- a/koperkapel/vehicles/base.py
+++ b/koperkapel/vehicles/base.py
@@ -1,5 +1,6 @@
 """ Base class for vehicles.  """
 
+from itertools import chain, islice, repeat
 from pygame.constants import BLEND_RGBA_MULT
 from pgzero.actor import Actor
 from pgzero.loaders import images
@@ -30,6 +31,22 @@ class Vehicle:
             for seat_pos, (roach, _) in roach_seating_numbers if roach
         }
 
+    def seat_roach(self, world, roach, seat_pos):
+        vehicle = world.vehicles[self.vehicle_type]
+        seats = len(self.seats)
+        seating = list(vehicle.seating)
+        seating = list(islice(
+            chain(seating, repeat(None, seats)), 0, seats))
+        seating[seat_pos] = roach
+        # line below records new seating on the world proxy
+        vehicle.seating = seating
+
+    def roach_at(self, world, seat_pos):
+        roach_seating = world.vehicles[self.vehicle_type].seating
+        if seat_pos >= len(roach_seating):
+            return None
+        return roach_seating[seat_pos]
+
     _vehicle_types = {}
 
     @classmethod