generators created and general cleanup

QOLfunc.py

@@ -14,3 +14,4 @@ def chooseNeighbours(count: int, i: int) -> set:
         if num != i:
             neigh.add(num)
     return neigh
+

classes/Household.py

@@ -8,33 +8,51 @@ class Household:
 
     def setNeighbours(self, neighbours):
         self.neighbours = neighbours
+        return self
 
     def addNeighbour(self, neighbour):
         self.neighbours.append(neighbour)
+        return self
 
     def getNeighbours(self):
         return self.neighbours
 
     def setPosition(self, position: tuple):
         self.position = position
+        return self
 
     def getPosition(self):
         return self.position
 
     def setImage(self, image: object):
         self.image = image
+        return self
 
     def getImage(self):
         return self.image
 
     def setId(self, newID: int):
         self.id = newID
+        return self
 
     def getId(self):
         return self.id
 
     def setGarbage(self, garbage: object):
         self.garbage = garbage
+        return self
 
     def getGarbage(self):
         return self.garbage
+
+
+def generateHousehold(i, image, position):
+    """
+    :param int i:
+    :param object image:
+    :param tuple position:
+    :return:
+    """
+    house = Household()
+    house.setId(i).setImage(image).setPosition(position)
+    return house

generators.py

@@ -0,0 +1,58 @@
+from classes.Junkyard import *
+from classes.Trashcan import *
+from classes.Household import *
+from QOLfunc import *
+import pygame
+
+
+def trashcanGenerator(junkyard: Junkyard) -> list:
+    typelist = ["paper", "metals_and_plastics", "mixed", "bio_waste", "glass"]
+    trash_cans = [
+        {'position': (1600, 300), 'type': 'paper'},
+        {'position': (1600, 500), 'type': 'metals_and_plastics'},
+        {'position': (1600, 700), 'type': 'mixed'},
+        {'position': (1600, 900), 'type': 'bio_waste'},
+        {'position': (1600, 100), 'type': 'glass'},
+    ]
+    trash_can_images = {
+        'paper': pygame.image.load('sprites/trash_can_papier.jpg').convert_alpha(),
+        'metals_and_plastics': pygame.image.load('sprites/trash_can_metale_plastik.jpg').convert_alpha(),
+        'mixed': pygame.image.load('sprites/trash_can_zmieszane.jpg').convert_alpha(),
+        'bio_waste': pygame.image.load('sprites/trash_can_bio.jpg').convert_alpha(),
+        'glass': pygame.image.load('sprites/trash_can_szklo.jpg').convert_alpha(),
+    }
+    trashcans = []
+    for key in trash_can_images:
+        trash_can_images[key] = pygame.transform.scale(trash_can_images[key], (90,90))
+    for i in range(5):
+        trashcan = generateTrashcan(i, trash_can_images[typelist[i]], trash_cans[i]["position"], trash_cans[i]["type"])
+        trashcan.addNeighbour(junkyard)
+        junkyard.addNeighbour(trashcan)
+        trashcans.append(trashcan)
+
+    return trashcans
+
+
+def householdGenerator(junkyard):
+    new_house_size = (140, 120)
+    house_positions = [(450, 100), (250, 250), (200, 700), (850, 800), (1000, 200), (1050, 700), (800, 100), (400, 850)]
+    houses = []
+    for i in range(8):
+        house_image = pygame.image.load(f'sprites/house_{i}.jpg')
+        house_image = pygame.transform.scale(house_image, new_house_size)
+        house = generateHousehold(i, house_image, house_positions[i])
+        houses.append(house)
+
+    for i in range(len(houses)):
+        if i % 2 == 0:
+            neigh = chooseNeighbours(2, i)
+        else:
+            neigh = chooseNeighbours(3, i)
+        for num in neigh:
+            if len(houses[i].getNeighbours()) < 3 and len(houses[num].getNeighbours()) < 3:
+                houses[i].addNeighbour(houses[num])
+                houses[num].addNeighbour(houses[i])
+        if i == 7 or i == 1:
+            houses[i].addNeighbour(junkyard)
+            junkyard.addNeighbour(houses[i])
+    return houses

main.py

@@ -1,8 +1,4 @@
-import pygame
-from classes.Trashcan import *
-from classes.Junkyard import *
-from classes.Household import *
-from QOLfunc import *
+from generators import *
 
 pygame.init()
 screen = pygame.display.set_mode((1800, 1000))
@@ -15,63 +11,20 @@ colors = {
 
 garbage_truck_image = pygame.image.load('sprites/garbage_truck.png')
 
-trash_can_images = {
-    'paper': pygame.image.load('sprites/trash_can_papier.jpg').convert_alpha(),
-    'metals_and_plastics': pygame.image.load('sprites/trash_can_metale_plastik.jpg').convert_alpha(),
-    'mixed': pygame.image.load('sprites/trash_can_zmieszane.jpg').convert_alpha(),
-    'bio_waste': pygame.image.load('sprites/trash_can_bio.jpg').convert_alpha(),
-    'glass': pygame.image.load('sprites/trash_can_szklo.jpg').convert_alpha(),
-}
-
 new_garbage_truck_size = (120, 120)
-new_trash_can_size = (90, 90)
 new_house_size = (140, 120)
 
 garbage_truck_image = garbage_truck_image.convert_alpha()
 garbage_truck_image = pygame.transform.scale(garbage_truck_image, new_garbage_truck_size)
 
-for key in trash_can_images:
-    trash_can_images[key] = pygame.transform.scale(trash_can_images[key], new_trash_can_size)
 
 garbage_truck_position = [800, 500]
 
-trash_cans = [
-    {'position': (1600, 300), 'type': 'paper'},
-    {'position': (1600, 500), 'type': 'metals_and_plastics'},
-    {'position': (1600, 700), 'type': 'mixed'},
-    {'position': (1600, 900), 'type': 'bio_waste'},
-    {'position': (1600, 100), 'type': 'glass'},
-]
+
 junkyard = Junkyard()
-trashcans = []
-typelist = ["paper", "metals_and_plastics", "mixed", "bio_waste", "glass"]
-for i in range(5):
-    trashcan = Trashcan(i, trash_can_images[typelist[i]], trash_cans[i]["position"], trash_cans[i]["type"])
-    trashcan.addNeighbour(junkyard)
-    junkyard.addNeighbour(trashcan)
-    trashcans.append(trashcan)
+trashcans = trashcanGenerator(junkyard)
+houses = householdGenerator(junkyard)
 
-houses = []
-house_positions = [(450, 100), (250, 250), (200, 700), (850, 800), (1000, 200), (1050, 700), (800, 100), (400, 850)]
-
-for i in range(8):
-    house_image = pygame.image.load(f'house_{i}.jpg')
-    house_image = pygame.transform.scale(house_image, new_house_size)
-    house = Household(i, house_image, house_positions[i])
-    houses.append(house)
-
-for i in range(len(houses)):
-    if i % 2 == 0:
-        neigh = chooseNeighbours(2, i)
-    else:
-        neigh = chooseNeighbours(3, i)
-    for num in neigh:
-        if len(houses[i].getNeighbours()) < 3 and len(houses[num].getNeighbours()) < 3:
-            houses[i].addNeighbour(houses[num])
-            houses[num].addNeighbour(houses[i])
-    if i == 7 or i == 1:
-        houses[i].addNeighbour(junkyard)
-        junkyard.addNeighbour(houses[i])
 
 running = True
 while running: