Change grid class to consist of nodes.

UI/grid.py

@@ -3,6 +3,29 @@ import numpy as np
 
 
 class Grid:
+
+    def __init__(self, cols: int, rows: int):
+        self.table = [[Node(row, col)
+                       for col in range(cols)]
+                      for row in range(rows)]
+        self.cols = cols
+        self.rows = rows
+
+    def draw_map(self, screen: "PyGame screen"):
+        """Draws whole map"""
+        screen.fill(Node.BLACK)
+        for row in self.table:
+            for node in row:
+                node.draw(screen)
+
+    def change_field(self, row: int, col: int, f_type: int):
+        self.table[row][col].change_field_type(f_type)
+
+    def draw_node(self, screen, row: int, col: int):
+        self.table[row][col].draw(screen)
+
+
+class Node:
     # define rectangles dimensions
     r_width = 20
     r_height = 20
@@ -16,32 +39,40 @@ class Grid:
     RED = (255, 0, 0)
     BLUE = (0, 0, 255)
 
-    def __init__(self, cols, rows):
+    def __init__(self, row: int, col: int,
-        self.table = np.zeros(shape=(rows, cols), dtype=int)
+                 field_type: int = 0, reachable: bool = True):
-        self.cols = list(range(cols))
+        self.row = row
-        self.rows = list(range(rows))
+        self.col = col
+        self.field_type = field_type
+        self.reachable = reachable
+        self.visited = False
 
-    def draw_map(self, screen):
+    def visit(self):
-        screen.fill(self.BLACK)
+        self.visited = True
-        for col in self.cols:
-            for row in self.rows:
-                self.draw_node(screen, row, col)
 
-    def draw_node(self, screen, row, col):
+    def draw(self, screen):
-        if self.table[row][col] == 0:
+        color = self.get_field_color()
-            color = self.GREEN
+        col = self.col
-        elif self.table[row][col] == 1:
+        row = self.row
-            color = self.RED
+        width = self.r_width
-        elif self.table[row][col] == 2:
+        height = self.r_height
-            color = self.BLUE
+        margin = self.r_margin
         # rect -> (left, top, width, height)
         # draw.rect(surface, color, rect, margin)
-
         pg.draw.rect(screen, color,
-                     ((col * (self.r_width + self.r_margin)) + self.r_margin,
+                     ((col * (width + margin)) + margin,
-                      (row * (self.r_height + self.r_margin)) + self.r_margin,
+                      (row * (height + margin)) + margin,
-                      self.r_width, self.r_height))
+                      width, height))
-        pg.display.flip()
+        pg.display.flip()  # refresh screen
 
-    def change_field(self, row, col, field_type):
+    def change_field_type(self, field_type: int):
-        self.table[row][col] = field_type
+        self.field_type = field_type
+
+    def get_field_color(self) -> tuple:
+        """Gets the color tuple of field"""
+        if self.field_type == 0:
+            return self.GREEN
+        elif self.field_type == 1:
+            return self.RED
+        elif self.field_type == 2:
+            return self.BLUE

UI/window.py

@@ -1,22 +1,22 @@
 import pygame as pg
 import time
 from pygame.locals import *
-from UI.grid import Grid
+from UI.grid import Grid, Node
 
 
 class Window():
-    def __init__(self, grid):
+    def __init__(self, grid: Grid):
         pg.init()
         # setup window
         pg.display.set_caption('Inteligentna śmieciarka')
 
         self.grid = grid
         # assign to variables for brevity
-        cols = len(self.grid.cols)
+        cols = self.grid.cols
-        rows = len(self.grid.rows)
+        rows = self.grid.rows
-        width = self.grid.r_width
+        width = Node.r_width
-        height = self.grid.r_height
+        height = Node.r_height
-        margin = self.grid.r_margin
+        margin = Node.r_margin
 
         screen_width = cols * (width + margin) + 2 * margin
         screen_height = rows * (height + margin) + 2 * margin