From d0053368e35b40313d1819d3007ccb2fd85ee784 Mon Sep 17 00:00:00 2001
From: Auric Vente <mprops@protonmail.com>
Date: Thu, 25 Jul 2024 09:09:42 -0600
Subject: [PATCH] Mods

---
 cromulant/game.py   | 21 ++++++++++++++-------
 cromulant/window.py | 31 ++++++++++++-------------------
 2 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/cromulant/game.py b/cromulant/game.py
index 5c6f961..e09e8a1 100644
--- a/cromulant/game.py
+++ b/cromulant/game.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import random
+from typing import Any
 
 from PySide6.QtWidgets import QHBoxLayout  # type: ignore
 from PySide6.QtWidgets import QVBoxLayout
@@ -8,6 +9,7 @@ from PySide6.QtWidgets import QLabel
 from PySide6.QtWidgets import QWidget
 from PySide6.QtWidgets import QFrame
 from PySide6.QtWidgets import QMenu
+from PySide6.QtWidgets import QDialog
 from PySide6.QtGui import QCursor  # type: ignore
 from PySide6.QtGui import QMouseEvent
 from PySide6.QtGui import QPixmap
@@ -22,6 +24,7 @@ from .utils import Utils
 from .ants import Ant
 from .ants import Ants
 from .window import Window
+from .window import RestartDialog
 from .settings import Settings
 
 
@@ -409,24 +412,28 @@ class Game:
 
     @staticmethod
     def restart() -> None:
-        opts = ["25", "50", "100", "250"]
+        sizes = ["25", "50", "100", "250"]
         defindex = 0
 
-        for i, opt in enumerate(opts):
+        for i, opt in enumerate(sizes):
             if int(opt) == Config.default_population:
                 defindex = i
                 break
 
-        opts = [f"{opt} ants" for opt in opts]
-        size = Window.prompt_combobox("Size of the population", opts, defindex)
+        size_opts = [f"{opt} ants" for opt in sizes]
+        dialog = RestartDialog(size_opts, defindex)
+        data: dict[str, Any] | None = None
 
-        if not size:
+        if dialog.exec() == QDialog.Accepted:
+            data = dialog.get_data()
+
+        if not data:
             return
 
-        num = int(size.split(" ")[0])
+        size = int(data["size"].split(" ")[0])
 
         Window.clear_view()
-        Ants.populate(num)
+        Ants.populate(size)
         Window.to_top()
         Game.intro()
         Game.start_loop()
diff --git a/cromulant/window.py b/cromulant/window.py
index 68be0c2..99c8198 100644
--- a/cromulant/window.py
+++ b/cromulant/window.py
@@ -39,21 +39,21 @@ class FilterLineEdit(QLineEdit):  # type: ignore
             super().keyPressEvent(e)
 
 
-class ComboBoxDialog(QDialog):  # type: ignore
-    def __init__(self, message: str, options: list[str], defindex: int) -> None:
+class RestartDialog(QDialog):  # type: ignore
+    def __init__(self, sizes: list[str], defindex: int) -> None:
         super().__init__()
         self.setWindowTitle("Select Option")
         self.setFixedSize(300, 150)  # Set a fixed size for the dialog
 
         self.layout = QVBoxLayout()
 
-        self.label = QLabel(message)
+        self.label = QLabel("Size of the population")
         self.layout.addWidget(self.label)
 
-        self.combo_box = QComboBox()
-        self.combo_box.addItems(options)
-        self.combo_box.setCurrentIndex(defindex)
-        self.layout.addWidget(self.combo_box)
+        self.size_combo = QComboBox()
+        self.size_combo.addItems(sizes)
+        self.size_combo.setCurrentIndex(defindex)
+        self.layout.addWidget(self.size_combo)
 
         self.button_layout = QHBoxLayout()
 
@@ -68,9 +68,12 @@ class ComboBoxDialog(QDialog):  # type: ignore
 
         self.layout.addLayout(self.button_layout)
         self.setLayout(self.layout)
+        self.setWindowFlags(Qt.Popup)
 
-    def get_selection(self) -> str:
-        return str(self.combo_box.currentText())
+    def get_data(self) -> dict[str, Any]:
+        return {
+            "size": str(self.size_combo.currentText()),
+        }
 
 
 class Window:
@@ -365,13 +368,3 @@ class Window:
         msg_box.setWindowTitle("Information")
         msg_box.setStandardButtons(QMessageBox.Ok)
         msg_box.exec()
-
-    @staticmethod
-    def prompt_combobox(message: str, options: list[str], defindex: int = 0) -> str:
-        dialog = ComboBoxDialog(message, options, defindex)
-        dialog.setWindowFlags(Qt.Popup)
-
-        if dialog.exec() == QDialog.Accepted:
-            return dialog.get_selection()
-
-        return ""