Coverage for /home/deng/Projects/metatree_drawer/metatreedrawer/treeprofiler/tree_image.py: 30%
37 statements
« prev ^ index » next coverage.py v7.2.7, created at 2024-08-07 10:33 +0200
« prev ^ index » next coverage.py v7.2.7, created at 2024-08-07 10:33 +0200
1from selenium import webdriver
2from selenium.webdriver.chrome.options import Options
3from selenium.webdriver.common.action_chains import ActionChains
4import time
5import os
6from ete4 import Tree, PhyloTree
7from multiprocessing import Process
9DRIVER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'chromedriver')
11def tree_session(tree, layouts, port):
12 #t = Tree(TREEFILE, format=1)
13 tree.explore(layouts=layouts, keep_server=True, open_browser=False, compress=False,
14 show_leaf_name=True, show_branch_length=True, show_branch_support=True, port=port,
15 )
16 #time.sleep(10)
18def snap_tree(port, plot_file):
19 url = 'http://127.0.0.1:{}/'.format(str(port))
21 options = Options()
22 options.headless = True
24 driver = webdriver.Chrome(executable_path=DRIVER, options=options)
25 driver.get(url)
26 time.sleep(0.5)
28 actions = ActionChains(driver)
29 actions.send_keys('d')
30 actions.perform()
32 defualt_download_dir = os.getcwd() # specify the directory where the file will be downloaded
33 file_name = 'example.svg' # specify the name of the downloaded file
34 file_path = os.path.join(defualt_download_dir, file_name)
36 # wait for the file to appear in the directory
37 while not os.path.exists(file_path):
38 time.sleep(1)
39 #file_path = 'examples/basic_example1/basic_example1_annotated.png'
41 # Move the file
42 os.rename(file_path, plot_file)
44 driver.quit()
46def get_image(tree, layouts, port, plot_file):
47 p1 = Process(target=tree_session, args=(tree, layouts, port,))
48 p2 = Process(target=snap_tree, args=(port, plot_file, ))
49 print(plot_file)
50 p1.start()
51 time.sleep(2)
52 p2.start()
53 time.sleep(4)
54 p1.kill()
55 p2.kill()
56 #p1.join()
57#main(annotated_tree, layouts, port, plot_file)