Docstrings

overcooked_simulator/gui_2d_vis/drawing.py

@@ -49,12 +49,26 @@ def create_polygon(n, length):
 
 
 class Visualizer:
+    """Class for visualizing the game state retrieved from the gameserver.
+    2D game screen is drawn with pygame shapes and images.
+
+    Args:
+        config: Visualization configuration (loaded from yaml file) given as a dict.
+
+    """
+
     def __init__(self, config):
         self.image_cache_dict = {}
         self.player_colors = []
         self.config = config
 
     def create_player_colors(self, n) -> None:
+        """Create different colors for the players. The color hues are sampled uniformly in HSV-Space,
+        then the corresponding colors from the defined colors list are looked up.
+
+        Args:
+            n: Number of players to create colors for.
+        """
         hue_values = np.linspace(0, 1, n + 1)
 
         colors_vec = np.array([col for col in colors.values()])
@@ -72,10 +86,18 @@ class Visualizer:
 
     def draw_gamescreen(
         self,
-        screen,
-        state,
-        grid_size,
+        screen: pygame.Surface,
+        state: dict,
+        grid_size: int,
     ):
+        """Draws the game state on the given surface.
+
+        Args:
+            screen: The pygame surface to draw the game on.
+            state: The gamestate retrieved from the environment.
+            grid_size: The gridsize to base every object size in the game on.
+        """
+
         width = int(np.ceil(state["kitchen"]["width"] * grid_size))
         height = int(np.ceil(state["kitchen"]["height"] * grid_size))
         self.draw_background(
@@ -96,8 +118,18 @@ class Visualizer:
             grid_size,
         )
 
-    def draw_background(self, surface, width, height, grid_size):
-        """Visualizes a game background."""
+    def draw_background(
+        self, surface: pygame.Surface, width: int, height: int, grid_size: int
+    ):
+        """Visualizes a game background.
+
+        Args:
+            surface: The pygame surface to draw the background on.
+            width: The kitchen width.
+            height: The kitchen height.
+            grid_size: The gridsize to base the background shapes on.
+        """
+
         block_size = grid_size // 2  # Set the size of the grid block
         surface.fill(colors[self.config["Kitchen"]["ground_tiles_color"]])
         for x in range(0, width, block_size):
@@ -118,6 +150,15 @@ class Visualizer:
         pos: npt.NDArray,
         rot_angle=0,
     ):
+        """Draws an image on the given screen.
+
+        Args:
+            screen: The pygame surface to draw the image on.
+            img_path: The path to the image file, given relative to the gui_2d_vis directory.
+            size: The size of the image, given in pixels.
+            pos: The position of the center of the image, given in pixels.
+            rot_angle: Optional angle to rotate the image around.
+        """
         cache_entry = f"{img_path}"
         if cache_entry in self.image_cache_dict.keys():
             image = self.image_cache_dict[cache_entry]
@@ -143,7 +184,11 @@ class Visualizer:
     ):
         """Visualizes the players as circles with a triangle for the facing direction.
         If the player holds something in their hands, it is displayed
-        Args:            state: The game state returned by the environment.
+
+        Args:
+            screen: The pygame surface to draw the players on.
+            players: The state of the players returned by the environment.
+            grid_size: The gridsize to rescale the drawn players to.
         """
         for p_idx, player_dict in enumerate(players):
             player_dict: PlayerState
@@ -371,7 +416,14 @@ class Visualizer:
         percent: float,
         grid_size: float,
     ):
-        """Visualize progress of progressing item as a green bar under the item."""
+        """Visualize progress of progressing item as a green bar under the item.
+
+        Args:
+            screen: The pygame surface to draw the progress bar on.
+            pos: The center position of a tile to draw the progress bar under.
+            percent: Progressed percent of the progress bar.
+            grid_size: Scaling of the progress bar given in pixels.
+        """
         pos -= grid_size / 2
 
         bar_height = grid_size * 0.2
@@ -390,7 +442,10 @@ class Visualizer:
         """Visualization of a counter at its position. If it is occupied by an item, it is also shown.
         The visual composition of the counter is read in from visualization.yaml file, where it is specified as
         different parts to be drawn.
-        Args:            counter: The counter to visualize.
+        Args:
+            screen: The pygame surface to draw the counter on.
+            counter_dict: The counter to visualize, given as a dict from the game state.
+            grid_size: Scaling of the counter given in pixels.
         """
         pos = np.array(counter_dict["pos"], dtype=float) * grid_size
         counter_type = counter_dict["type"]
@@ -435,6 +490,14 @@ class Visualizer:
         pos: npt.NDArray[float],
         item_scale: float,
     ):
+        """Visualization of a thing lying on a counter.
+        Args:
+            screen: The pygame surface to draw the item on the counter on.
+            occupied_by: The thing that occupies the counter.
+            grid_size: Scaling of the object given in pixels.
+            pos: The position of the counter which the thing lies on.
+            item_scale: Relative scaling of the item.
+        """
         # Multiple plates on plate return:
         if isinstance(occupied_by, list):
             for i, o in enumerate(occupied_by):
@@ -455,10 +518,13 @@ class Visualizer:
                 scale=item_scale,
             )
 
-    def draw_counters(self, screen: pygame, counters, grid_size):
+    def draw_counters(self, screen: pygame, counters: dict, grid_size: int):
         """Visualizes the counters in the environment.
 
-        Args:            state: The game state returned by the environment.
+        Args:
+            screen: The pygame surface to draw the counters on.
+            counters: The counter state returned by the environment.
+            grid_size: Scaling of the object given in pixels.
         """
         for counter in counters:
             self.draw_counter(screen, counter, grid_size)
@@ -514,8 +580,27 @@ class Visualizer:
                 )
 
     def draw_orders(
-        self, screen, state, grid_size, width, height, screen_margin, config
+        self,
+        screen: pygame.surface,
+        state: dict,
+        grid_size: int,
+        width: int,
+        height: int,
+        screen_margin: int,
+        config: dict,
     ):
+        """Visualization of the current orders.
+
+        Args:
+            screen: pygame surface to draw the orders on, probably not the game screen itself.
+            state: The game state returned by the environment.
+            grid_size: Scaling of the drawn orders, given in pixels.
+            width: Width of the pygame window
+            height: Height of the pygame window.
+            screen_margin: Size of the space around the game screen, for buttons, ... .
+            config: Visualization configuration (loaded from yaml file) given as a dict.
+
+        """
         orders_width = width - 100
         orders_height = screen_margin
         order_screen = pygame.Surface(
@@ -582,6 +667,14 @@ class Visualizer:
     def save_state_image(
         self, grid_size: int, state: dict, filename: str | Path
     ) -> None:
+        """Saves a screenshot of the visualization of the given state.
+
+        Args:
+            grid_size: Scaling of the world elements given in pixels.
+            state: Game state returned by the environment.
+            filename: Filename to save the screenshot to.
+
+        """
         width = int(np.ceil(state["kitchen"]["width"] * grid_size))
         height = int(np.ceil(state["kitchen"]["height"] * grid_size))
 
@@ -593,6 +686,15 @@ class Visualizer:
 
 
 def save_screenshot(state: dict, config: dict, filename: str | Path) -> None:
+    """Standalone function to save a screenshot. Creates a visualizer from the config and visualizes
+    the game state, saves it to the given filename.
+
+    Args:
+        state: The gamestate to visualize.
+        config: Visualization config for the visualizer.
+        filename: Filename to save the image to.
+
+    """
     viz = Visualizer(config)
     viz.create_player_colors(len(state["players"]))
     pygame.init()