simplified code

screenshot.py

@@ -14,6 +14,7 @@
 
 
 class ScreenshotHandler:
+    """Handles the screenshotting and sending of the screenshot."""
     SCREENSHOT_PATH = os.path.join(os.path.dirname(__file__), "Screenshots")
     MAX_NUMBER_OF_FILES = (
         1000 + 1
@@ -24,35 +25,32 @@ def __init__(self, navi, interaction, original_message):
         self.interaction = interaction
         self.original_message = original_message
         self.view = None
+        self.pil_image = None
 
     async def save_send_screenshot(self):
         """saves a screenshot of the current page and sends it to the user."""
         await self.navi.wait_for_not_load_screen()
 
-        file_number = self.get_next_file_number(self.SCREENSHOT_PATH)
-        new_screenshot_path = os.path.join(self.SCREENSHOT_PATH, f"{file_number}.png")
-
         (
             selector,
             crop,
             self.view,
-            new_screenshot_path,
             b64_image
-        ) = await self.get_screenshot_info_by_stage(new_screenshot_path)
+        ) = await self.get_screenshot_info_by_stage()
 
         if b64_image:  # for final stage, where the image is a base64 string, not a screenshot of an element.
             image_bytes = b64.b64decode(b64_image)
-            pil_image = Image.open(BytesIO(image_bytes))
+            self.pil_image = Image.open(BytesIO(image_bytes))
         else:
-            pil_image = await self.navi.screenshot(selector)
+            self.pil_image = await self.navi.screenshot(selector)
 
         if crop:
-            width, height = pil_image.size
-            pil_image = pil_image.crop((0, height - 630, width, height))
+            width, height = self.pil_image.size
+            self.pil_image = self.pil_image.crop((0, height - 630, width, height))
 
         # Convert the PIL image to bytes
         byte_arr = BytesIO()
-        pil_image.save(byte_arr, format="PNG")
+        self.pil_image.save(byte_arr, format="PNG")
         byte_arr.seek(0)
 
         file = nextcord.File(byte_arr, filename="image.png")
@@ -62,7 +60,7 @@ async def save_send_screenshot(self):
         )
 
         # Close the PIL image
-        del pil_image
+        self.pil_image = None
         byte_arr.close()
 
         self.original_message = await self.original_message.fetch()
@@ -72,13 +70,6 @@ async def save_send_screenshot(self):
             await self.view.wait()
             await self.add_reaction()
 
-    def get_next_file_number(self, path):
-        """Get the next available file number in the given directory."""
-        existing_files = glob.glob(os.path.join(path, "*.png"))
-        existing_numbers = {int(os.path.splitext(os.path.basename(f))[0]) for f in existing_files}
-        file_number = next(i for i in range(self.MAX_NUMBER_OF_FILES) if i not in existing_numbers)
-        return file_number
-
     async def remove_reaction(self):
         """Removes all reactions from the original message."""
         if (
@@ -107,42 +98,37 @@ async def add_reaction(self):
                 asyncio.create_task(self.original_message.add_reaction(self.view.current_label))
             asyncio.create_task(self.original_message.add_reaction("⏳"))
 
-    async def get_screenshot_info_by_stage(self, new_screenshot_path):
-        """Returns the selector, crop, view, and new_screenshot_path based on the stage of the grid, as well as the base64 image if on the last stage."""
-        selector, crop, view, b64_image = (
-            None,
-            False,
-            None,
-            None,
-        )  # placeholder values.
+    async def get_screenshot_info_by_stage(self):
+        """Returns the selector, crop, and view based on the stage of the grid, as well as the base64 image if on the last stage."""
 
-        if View.stage[self.interaction.user.id] in range(1, 4):
+        stage = View.stage[self.interaction.user.id]
+
+        if stage in range(1, 4):
             selector = ".waifu-container"
             crop = True
             view = View(self.navi, self.interaction)
+            b64_image = None
 
-        elif View.stage[self.interaction.user.id] == 0:
+        elif stage == 0:
             selector = ".waifu-grid"
+            crop = False
             view = View(self.navi, self.interaction)
+            b64_image = None
 
         else:
             # if on last stage.
             # this is a bit of a hack, but it works
             # the final image is not a screenshot, but a base64 image taken from the element's src attribute
             # this is done because of an issue with pyppeteer's screenshot func
-            file_number = self.get_next_file_number(os.path.join(self.SCREENSHOT_PATH, "final_results"))
-            new_screenshot_path = os.path.join(
-                self.SCREENSHOT_PATH, "final_results", f"{file_number}.png"
-            )
             await self.navi.wait_for_final_image()
 
-            final_image_element = await self.navi.page.querySelector(
-                ".waifu-preview > img"
-            )
-            final_image_url = await self.navi.page.evaluate(
-                "(element) => element.src", final_image_element
-            )
+            final_image_element = await self.navi.page.querySelector(".waifu-preview > img")
+            final_image_url = await self.navi.page.evaluate("(element) => element.src", final_image_element)
             final_image_base64 = final_image_url.split(",")[1]
+
+            selector = None
+            crop = False
+            view = None
             b64_image = final_image_base64
 
-        return selector, crop, view, new_screenshot_path, b64_image
+        return selector, crop, view, b64_image