add colors (#4)

HokageM · Feb 24, 2024 · 2ebeb54 · 2ebeb54
1 parent dc203f7
commit 2ebeb54
Showing 1 changed file with 16 additions and 7 deletions.
diff --git a/src/ninjabees/environment/World.py b/src/ninjabees/environment/World.py
@@ -8,6 +8,13 @@ class World:
     This class represents the world of the bee simulation.
     """
 
+    RED = '\033[0;31m'
+    BLUE = '\033[0;34m'
+    GREEN = '\033[0;32m'
+    YELLOW = '\033[0;33m'
+    MAGENTA = '\033[0;35m'
+    RESET = '\033[0m'
+
     def __init__(self, width, height):
         self.__width = width
         self.__height = height
@@ -19,7 +26,7 @@ def __init__(self, width, height):
         self.__unclaimed_food_sources = []
         self.__n_food_sources = 0
 
-        self.__world_map = [['-' for _ in range(width)] for _ in range(height)]
+        self.__world_map = [[' ' for _ in range(width)] for _ in range(height)]
 
     def add_food_source(self, food_source):
         """
@@ -71,20 +78,22 @@ def update_world_map(self):
         :return:
         """
         for food in self.__unclaimed_food_sources:
-            self.__world_map[food.get_y()][food.get_x()] = 'u'
+            self.__world_map[food.get_y()][food.get_x()] = f'{World.RED}u{World.RESET}'
         for entity in self.__entities:
             entity_type = entity.get_type()
             entity_x = entity.get_x()
             entity_y = entity.get_y()
             if entity_type == EntityType.Bee:
                 if entity.has_found_food():
-                    self.__world_map[entity_y][entity_x] = 'S' if entity.get_job() == BeeJob.Scout else 'B'
+                    self.__world_map[entity_y][entity_x] = f'{World.BLUE}S{World.RESET}' \
+                        if entity.get_job() == BeeJob.Scout else f'{World.YELLOW}B{World.RESET}'
                 else:
-                    self.__world_map[entity_y][entity_x] = 's' if entity.get_job() == BeeJob.Scout else 'b'
+                    self.__world_map[entity_y][entity_x] = f'{World.BLUE}s{World.RESET}' \
+                        if entity.get_job() == BeeJob.Scout else f'{World.YELLOW}b{World.RESET}'
             elif entity_type == EntityType.Food:
-                self.__world_map[entity_y][entity_x] = 'F'
+                self.__world_map[entity_y][entity_x] = f'{World.GREEN}F{World.RESET}'
             elif entity_type == EntityType.Hive:
-                self.__world_map[entity_y][entity_x] = 'H'
+                self.__world_map[entity_y][entity_x] = f'{World.MAGENTA}H{World.RESET}'
             else:
                 raise ValueError("Unknown entity type")
 
@@ -101,7 +110,7 @@ def run(self, max_iterations):
             Animator.print_world_status(world_map=self.__world_map, found=len(self.__hive.found_food_sources),
                                         total=self.__n_food_sources, food_at_hive=self.__hive.food_at_hive)
 
-            self.__world_map = [['-' for _ in range(self.__width)] for _ in range(self.__height)]
+            self.__world_map = [[' ' for _ in range(self.__width)] for _ in range(self.__height)]
 
             if len(self.__hive.found_food_sources) == self.__n_food_sources:
                 print(f'All food sources found! In {iteration} iterations')