cellular_automata/wa_tor.py

@@ -15,7 +15,7 @@ to result in one of the three possible results.
 from collections.abc import Callable
 from random import randint, shuffle
 from time import sleep
-from typing import Literal
+from typing import Any, Literal
 
 WIDTH = 50  # Width of the Wa-Tor planet
 HEIGHT = 50  # Height of the Wa-Tor planet
@@ -54,7 +54,7 @@ class Entity:
         self.coords = coords
 
         self.remaining_reproduction_time = (
-            PREY_REPRODUCTION_TIME if prey else PREDATOR_REPRODUCTION_TIME
+            PREY_REPRODUCTION_TIME if prey is True else PREDATOR_REPRODUCTION_TIME
         )
         self.energy_value = None if prey is True else PREDATOR_INITIAL_ENERGY_VALUE
         self.alive = True
@@ -124,9 +124,9 @@ class WaTor:
 
         # Populate planet with predators and prey randomly
         for _ in range(PREY_INITIAL_COUNT):
-            self.add_entity(prey=True)
+            self.add_entity(True)
         for _ in range(PREDATOR_INITIAL_COUNT):
-            self.add_entity(prey=False)
+            self.add_entity(False)
         self.set_planet(self.planet)
 
     def set_planet(self, planet: list[list[Entity | None]]) -> None:
@@ -167,8 +167,8 @@ class WaTor:
         while True:
             row, col = randint(0, self.height - 1), randint(0, self.width - 1)
             if self.planet[row][col] is None:
-                self.planet[row][col] = Entity(prey=prey, coords=(row, col))
-                return
+                break
+        self.planet[row][col] = Entity(prey=prey, coords=(row, col))
 
     def get_entities(self) -> list[Entity]:
         """
@@ -178,7 +178,11 @@ class WaTor:
         >>> len(wt.get_entities()) == PREDATOR_INITIAL_COUNT + PREY_INITIAL_COUNT
         True
         """
-        return [entity for column in self.planet for entity in column if entity]
+        start: Any = []
+        return sum(
+            [[entity for entity in column if entity] for column in self.planet],
+            start=start,
+        )
 
     def balance_predators_and_prey(self) -> None:
         """
@@ -199,18 +203,19 @@ class WaTor:
         shuffle(entities)
 
         if len(entities) >= MAX_ENTITIES - MAX_ENTITIES / 10:
-            prey = [entity for entity in entities if entity.prey]
-            predators = [entity for entity in entities if not entity.prey]
+            prey = list(filter(lambda entity: entity.prey is True, entities))
+            predators = list(filter(lambda entity: entity.prey is True, entities))
 
             prey_count, predator_count = len(prey), len(predators)
 
-            entities_to_purge = (
-                prey[:DELETE_UNBALANCED_ENTITIES]
-                if prey_count > predator_count
-                else predators[:DELETE_UNBALANCED_ENTITIES]
-            )
-            for entity in entities_to_purge:
-                self.planet[entity.coords[0]][entity.coords[1]] = None
+            if prey_count > predator_count:
+                for entity in prey[:DELETE_UNBALANCED_ENTITIES]:
+                    # Purge the first n entities of the prey
+                    self.planet[entity.coords[0]][entity.coords[1]] = None
+            else:
+                for entity in predators[:DELETE_UNBALANCED_ENTITIES]:
+                    # Purge the first n entities of the predators
+                    self.planet[entity.coords[0]][entity.coords[1]] = None
 
     def get_surrounding_prey(self, entity: Entity) -> list[Entity]:
         """
@@ -330,7 +335,8 @@ class WaTor:
         Entity(prey=False, coords=(0, 1), remaining_reproduction_time=20,
         energy_value=15)]]
         """
-        row, col = coords = entity.coords
+        coords = entity.coords
+        row, col = coords
 
         for direction in direction_orders:
             # If the direction is North and the northern square
@@ -509,7 +515,7 @@ class WaTor:
                 self.time_passed(self, iter_num)
 
 
-def visualise(wt: WaTor, iter_number: int, *, colour: bool = True) -> None:
+def display_visually(wt: WaTor, iter_number: int, *, colour: bool = True) -> None:
     """
     Visually displays the Wa-Tor planet using
     an ascii code in terminal to clear and re-print
@@ -527,27 +533,27 @@ def visualise(wt: WaTor, iter_number: int, *, colour: bool = True) -> None:
     ... [Entity(False, coords=(1, 0)), None, Entity(False, coords=(1, 2))],
     ... [None, Entity(True, coords=(2, 1)), None]
     ... ])
-    >>> visualise(wt, 0, colour=False)  # doctest: +NORMALIZE_WHITESPACE
+    >>> display_visually(wt, 0, colour=False)  # doctest: +NORMALIZE_WHITESPACE
     #  x  .
     x  .  x
     .  #  .
     <BLANKLINE>
     Iteration: 0 | Prey count: 2 | Predator count: 3 |
     """
-    if colour:
+    if colour is True:
         __import__("os").system("")
         print("\x1b[0;0H\x1b[2J\x1b[?25l")
 
-    reprint = "\x1b[0;0H" if colour else ""
-    ansi_colour_end = "\x1b[0m " if colour else " "
+    reprint = "\x1b[0;0H" if colour is True else ""
+    ansii_colour_end = "\x1b[0m " if colour is True else " "
 
     planet = wt.planet
     output = ""
 
     # Iterate over every entity in the planet
-    for row in planet:
-        for entity in row:
-            if entity is None:
+    for i in range(len(planet)):
+        for j in range(len(planet[0])):
+            if (entity := planet[i][j]) is None:
                 output += " . "
             else:
                 if colour is True:
@@ -556,12 +562,12 @@ def visualise(wt: WaTor, iter_number: int, *, colour: bool = True) -> None:
                         if entity.prey
                         else "\x1b[38;2;255;255;15m"
                     )
-                output += f" {'#' if entity.prey else 'x'}{ansi_colour_end}"
+                output += f" {'#' if entity.prey else 'x'}{ansii_colour_end}"
 
         output += "\n"
 
     entities = wt.get_entities()
-    prey_count = sum(entity.prey for entity in entities)
+    prey_count = len(list(filter(lambda entity: entity.prey is True, entities)))
 
     print(
         f"{output}\n Iteration: {iter_number} | Prey count: {prey_count} | "
@@ -577,5 +583,5 @@ if __name__ == "__main__":
     doctest.testmod()
 
     wt = WaTor(WIDTH, HEIGHT)
-    wt.time_passed = visualise
+    wt.time_passed = display_visually
     wt.run(iteration_count=100_000)