Text-to-Chart / screenshot.py
Omnibus's picture
Update screenshot.py
36d0d44 verified
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
import urllib
#from selenium_webdriver import WebElement
import uuid
driver_type = 'chromedriver'
driver=False
def create_ss(graph,width=1000,height=8000):
uid=uuid.uuid4()
code_out=""
for ea in graph.split("\n"):
code=ea.strip().strip("\n")
code_out+=code
url=f"https://omnibus-mermaid-script.static.hf.space/index.html?mermaid={urllib.parse.quote_plus(code_out)}"
is_url=True
wait_time=10000
if is_url:
options = webdriver.ChromeOptions()
options.add_argument('--headless')
options.add_argument('--no-sandbox')
options.add_argument('--disable-dev-shm-usage')
mes='<center>operation success'
try:
driver = webdriver.Chrome(options=options)
driver.get(url)
html=driver.page_source
print(html)
driver.implicitly_wait(int(wait_time*30))
driver.set_window_size(int(width), int(height))
obj = driver.find_element(By.ID, "chart")
screenshot = obj.screenshot(f'{uid}-tmp.png')
return(f'{uid}-tmp.png')
except Exception as e:
print(e)
return None