importing functions from autobot instead of duplication

Author: emmanuel

.gitignore

@@ -1,4 +1,4 @@
-data/
+data/*
 .idea/
 treetime/*
 scripts/geckodriver.log

scripts/ChunkyBot.py

@@ -12,155 +12,11 @@
 import getpass
 
 from gotoh2 import *
+from autobot import get_driver, login, retrieve_genomes
 
 todaystr = datetime.strftime(datetime.now(), '%Y-%m-%d')
 cwd = os.getcwd()
 
-def get_driver(download_folder, executable_path):
-	"""
-	Instantiate remote control interface for Firefox web browser
-	
-	:param download_folder:  path to write downloaded files
-	:param executable_path:  path to geckodriver executable
-	:return:
-	"""
-	profile = webdriver.FirefoxProfile()
-	profile.set_preference('browser.download.folderList', 2)
-	profile.set_preference('browser.download.manager.showWhenStarting', False)
-	profile.set_preference("browser.download.dir", download_folder)
-	profile.set_preference('browser.helperApps.alwaysAsk.force', False)
-	profile.set_preference("browser.helperApps.neverAsk.saveToDisk", "application/octet-stream,octet-stream")
-	profile.set_preference("browser.helperApps.neverAsk.openFile", "application/octet-stream,octet-stream")
-
-	opts = Options()
-	opts.headless = True  # opts.set_headless()
-	assert opts.headless
-
-	return webdriver.Firefox(firefox_profile=profile, options=opts, executable_path=executable_path)
-
-
-def login(driver):
-	"""
-	Use GISAID access credentials to login to database.
-	:param driver: webdriver.Firefox object
-	:return:
-	"""
-	driver.get('https://www.epicov.org/epi3/cfrontend')
-	time.sleep(15)  # seconds
-
-	# read login credentials from Environment Variables
-	try:
-		user = os.environ['gisaid_u_variable']
-		pw = os.environ['gisaid_pw_variable']
-	except KeyError:
-		# variables not set, get access credentials interactively
-		user = getpass.getpass(prompt='GISAID username: ')
-		pw = getpass.getpass(prompt='GISAID password: ')
-
-	print('logging in')
-	driver.execute_script('document.getElementById("elogin").value="{}"'.format(user))
-	driver.execute_script('document.getElementById("epassword").value="{}"'.format(pw))
-	time.sleep(5)
-
-	# call javascript login function
-	driver.execute_script('doLogin()')
-	time.sleep(5)
-
-	#navigate to corona virus page
-	print('navigating to CoV db')
-	element = driver.find_element_by_xpath("//*[contains(text(), 'Browse')]")
-	element.click()
-	time.sleep(5)
-
-	return driver
-
-def find_prefix(driver):
-	#find prefix variable
-	element = driver.find_element_by_xpath("//div[@class='buttons container-slot']")
-	return element.get_attribute('id').split('_')[1]
-
-def retrieve_genomes(driver, start, end, download_folder):
-	"""
-	Retrieve genomes with a specified deposition date range in the GISAID database.
-	Adding several time delays to avoid spamming the database.
-
-	:param driver:  webdriver.Firefox object from login()
-	:param start:  date in ISO format (yyyy-mm-dd)
-	:param end:  date in ISO format (yyyy-mm-dd)
-	:return:  path to file download
-	"""
-
-	# find prefix variable
-	element = driver.find_element_by_xpath("//div[@class='buttons container-slot']")
-	htmlid_as_list = element.get_attribute('id').split('_')
-	variable = htmlid_as_list[1]
-
-	# navigate to corona virus page
-	print('navigating to CoV db')
-	element = driver.find_element_by_xpath("//button[contains(text(), 'Reset')]")
-	element.click()
-	time.sleep(5)
-
-	# trigger selection change
-	time_string = '[id^="ce_' + variable + '"][id$="_input"]'
-
-	driver.execute_script("document.querySelectorAll('{}')[2].value = '{}'".format(time_string, start))
-	driver.execute_script("document.querySelectorAll('{}')[2].onchange()".format(time_string))
-
-	driver.execute_script("document.querySelectorAll('{}')[3].value = '{}'".format(time_string, end))
-	driver.execute_script("document.querySelectorAll('{}')[3].onchange()".format(time_string))
-
-	driver.execute_script("document.querySelectorAll('[id^=\"ce_{}\"][id$=_input]')[2].onchange()".format(variable))
-	time.sleep(15)
-
-	print('selecting all seqs')
-	element = driver.find_element_by_xpath("//*[contains(text(), 'Total')]")
-	count = element.get_attribute('innerHTML').split()[1].replace(',', '')
-	if int(count) > 10000:
-		time.sleep(15)
-
-	checkbox = driver.find_element_by_xpath("//span[@class='yui-dt-label']/input[@type='checkbox']")
-	checkbox.click()
-	time.sleep(5)
-
-	# download seqs
-	element = driver.find_element_by_xpath("//*[contains(text(), 'Download')]")
-	driver.execute_script("arguments[0].click();", element)
-	time.sleep(5)
-
-	# switch to iframe to download
-	driver.switch_to_frame(driver.find_element_by_tag_name("iframe"))
-	print("Download")
-	time.sleep(5)
-
-	button = driver.find_element_by_xpath(
-		"//*[contains(text(), 'Download')]//ancestor::div[@style='float: right']"
-	)
-	button.click()
-	time.sleep(5)
-
-	# wait for download to complete
-	while True:
-		files = os.listdir(download_folder)
-		if len(files) == 0:
-			# download has not started yet
-			time.sleep(10)
-			continue
-		if any([f.endswith('.part') for f in files]):
-			time.sleep(5)
-			continue
-		break
-
-	print('Downloading complete')
-	downloaded_file = os.listdir(download_folder)[0]
-	driver.switch_to.default_content()
-
-	# reset browser
-	element = driver.find_element_by_xpath("//button[contains(text(), 'Reset')]")
-	element.click()
-	return os.path.join(download_folder, downloaded_file)
-
-
 def compare_dicts(old_fasta, new_fasta):
 	old_dict = dict(convert_fasta(open(old_fasta)))
 	new_dict = dict(convert_fasta(open(new_fasta)))
@@ -174,6 +30,8 @@ def compare_dicts(old_fasta, new_fasta):
 			block_diff[key] = value
 	return block_diff
 
+
+
 def parse_args():
 	""" Command line interface """
 	parser = argparse.ArgumentParser(

scripts/autobot.py

@@ -62,6 +62,12 @@ def login(driver):
 	driver.execute_script('doLogin()')
 	time.sleep(5)
 
+        # navigate to corona virus page
+	print('navigating to CoV db')
+	element = driver.find_element_by_xpath("//*[contains(text(), 'Browse')]")
+	element.click()
+	time.sleep(5)
+
 	return driver
 
 
@@ -81,12 +87,6 @@ def retrieve_genomes(driver, start, end, download_folder):
 	htmlid_as_list = element.get_attribute('id').split('_')
 	variable = htmlid_as_list[1]
 
-	# navigate to corona virus page
-	print('navigating to CoV db')
-	element = driver.find_element_by_xpath("//*[contains(text(), 'Browse')]")
-	element.click()
-	time.sleep(5)
-
 	# trigger selection change
 	time_string = '[id^="ce_' + variable + '"][id$="_input"]'
 
@@ -133,7 +133,6 @@ def retrieve_genomes(driver, start, end, download_folder):
 			time.sleep(10)
 			continue
 		if any([f.endswith('.part') for f in files]):
-			# FIXME: is this platform specific?
 			time.sleep(5)
 			continue
 		break
@@ -142,6 +141,7 @@ def retrieve_genomes(driver, start, end, download_folder):
 	downloaded_file = os.listdir(download_folder)[0]
 
 	# reset browser
+	time.sleep(30)
 	driver.switch_to.default_content()
 	element = driver.find_element_by_xpath("//button[@class='sys-event-hook sys-form-button']")
 	element.click()