import os import sys import datetime import subprocess from time import sleep from selenium import webdriver from selenium.webdriver.common.by import By from selenium.webdriver.support.ui import WebDriverWait from selenium.webdriver.support import expected_conditions as EC from PIL import Image from selenium import webdriver from PIL import Image import shutil from screenshotone import Client, TakeOptions from pathlib import Path from dotenv import load_dotenv load_dotenv(Path('../.env')) from loguru import logger logger.remove() logger.add(sys.stdout, colorize=True, format="{level}|{file}:{line}:{function}| {message}" ) def screenshot_api(url, generated_image_filepath): """ Use screenshotone API to take company webpage screenshots """ try: # create API client client = Client(os.getenv('SCREENSHOTONE_ACCESS_KEY'), os.getenv('SCREENSHOTONE_SECRET_KEY')) # set up options options = (TakeOptions.url(url) .format("png") .viewport_width(1024) .viewport_height(768) .block_cookie_banners(True) .block_chats(True)) # generate the screenshot URL and share it with a user #url = client.generate_take_url(options) # or render a screenshot and download the image as stream image = client.take(options) # store the screenshot the example.png file with open(generated_image_filepath, 'wb', encoding="utf-8") as result_file: shutil.copyfileobj(image, result_file) # Display the screenshot using Image.show image = Image.open(generated_image_filepath) image.show() # Wait for 2 seconds (adjust the delay as needed) sleep(2) # Close the image window image.close() except Exception as err: print(f"Failed in screenshotone api: {err}") generated_image_filepath = take_screenshot(url, generated_image_filepath) return generated_image_filepath def take_screenshot(url, generated_image_filepath): # Create a webdriver instance in headless mode options = webdriver.ChromeOptions() options.add_argument("--headless") driver = webdriver.Chrome(options=options) logger.debug(f"Taking screenshot of url: {url}") try: # Navigate to the given url driver.get(url) # Optionally, increase the delay to ensure all content is loaded sleep(2) # Explicitly wait for the page to load (adjust timeout as needed) WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.TAG_NAME, "body"))) # Set a larger window size driver.set_window_size(1200, 800) # Take a screenshot of the webpage screenshot = driver.get_screenshot_as_png() # Save the screenshot to a file with open(generated_image_filepath, "wb", encoding="utf-8") as f: f.write(screenshot) # Display the screenshot using Image.show image = Image.open(generated_image_filepath) image.show() # Wait for 2 seconds (adjust the delay as needed) sleep(2) # Close the image window using subprocess (platform-dependent) subprocess.run(["pkill", "-f", "display"]) # Adjust based on your platform and viewer # If using macOS, you can use the following: # subprocess.run(["osascript", "-e", 'tell application "Preview" to close every window']) # If using Windows, you can use the following: # subprocess.run(["taskkill", "/F", "/IM", "Microsoft.Photos.exe"]) logger.debug(f"Screenshot successfully stored at: {generated_image_filepath}") return generated_image_filepath finally: # Close the webdriver instance driver.quit()