Files
ALwrity/lib/utils/take_url_screenshot.py
2024-02-05 15:15:07 +05:30

114 lines
3.7 KiB
Python

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>{level}</level>|<green>{file}:{line}:{function}</green>| {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') 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") 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()