#!/usr/bin/env python3 import os import sys import time import getpass import webbrowser from optparse import OptionParser, SUPPRESS_HELP # You may have to install the requests module from pypy import requests import PyZenity # Specify some basic information about our command usage = "usage: %prog [arguments]" parser = OptionParser(usage=usage, version="%prog .4", description="Monitor a CS-Rosetta job or submit a new CS-Rosetta job. Only academic use is " "allowed.") # Specify the common arguments parser.add_option("--key", action="store", dest="key", default=None, help="Monitor an entry by key as it moves through CS-Rosetta rather than submit an entry.") parser.add_option("--interactive", action="store_true", dest="interactive", default=False, help="Interactively query for the information needed to submit.") parser.add_option("--email", action="store", dest="auth_email", default="invalid_email", help="Your e-mail address.") parser.add_option("--structures", action="store", dest="structures", default="1500", help="How many structures to submit.") parser.add_option("--protein-name", action="store", dest="prot_name", default="Automated submission", help="Protein name.") parser.add_option("--author-name", action="store", dest="auth_name", default="Automated submission", help="Name of human submitter.") parser.add_option("--trimming-method", action="store", type="choice", choices=["lange", "yang", "notrim"], dest="trim_type", default="lange", help="How to handle flexible tails.") parser.add_option("--cs-type", action="store", type="choice", choices=["star", "talos"], dest="chemical_shift_type", default="talos", help="What format is the chemical shift file? (star/talos)") parser.add_option("--star-file", action="store", type="str", dest="star_file", default=None, help="The location of a NMR-STAR file to submit. Or submit a TALOS file.") parser.add_option("--talos-file", action="store", type="str", dest="talos_file", default=None, help="The location of a TALOS file to submit. Or submit a STAR file.") parser.add_option("--rdc-file", action="store", type="str", dest="rdc_file", help="The location of an RDC file to submit.") parser.add_option("--constraint-file", action="store", type="str", dest="constraint_file", help="The location of a rosetta contraints format file to submit.") parser.add_option("--disulfide-file", action="store", type="str", dest="disulfide_file", help="The location of a disulfide linkages file to submit.") parser.add_option("--directory-mode", action="store", type="str", dest="directory_mode", default=None, help="Specify a directory. All sub directories will be scanned and all combinations of " "chemical shifts and constraints files in those folders will be submitted.") # Don't show the server URL option normally parser.add_option("--server", action="store", type="str", dest="server_url", default="https://csrosetta.bmrb.io", help=SUPPRESS_HELP) # Options, parse 'em (options, args) = parser.parse_args() def print_job(dictionary): job_dir = dictionary.pop('directory') full_path = dictionary.pop('path') dictionary['name'] = "AUTO from " + job_dir +\ " (" + "|".join([os.path.basename(x) for x in dictionary.values()]) + ")" print("Job name: %s" % dictionary['name']) if 'star' in dictionary: print("\tSTAR file: %s" % dictionary['star']) if 'talos' in dictionary: print("\tTalos file: %s" % dictionary['talos']) if 'rdc_file' in dictionary: print("\tRDC file: %s" % dictionary['rdc_file']) if 'disulfide_file' in dictionary: print("\tDisulfide bond file: %s" % dictionary['disulfide_file']) if 'constraints_file' in dictionary: print("\tConstraints file: %s" % dictionary['constraints_file']) dictionary['path'] = full_path print("\n") return dictionary def scan_folder_for_jobs(to_scan): from copy import deepcopy submissions = [] for item in os.listdir(to_scan): file_path = os.path.join(to_scan, item) if os.path.isdir(file_path): fnew = {'directory': item, 'path': file_path} files_in_dir = os.listdir(file_path) # Get the files for each_file in files_in_dir: ffpath = os.path.join(file_path, each_file) if each_file.endswith(".str"): fnew['star'] = ffpath if each_file.endswith(".tab"): fnew['talos'] = ffpath if each_file.endswith(".dsb"): fnew['disulfide_file'] = ffpath if each_file.endswith(".rdc"): fnew['rdc_file'] = ffpath # Make a submission for each constraint file for each_file in files_in_dir: if each_file.endswith(".cst"): copy = deepcopy(fnew) copy['constraints_file'] = os.path.join(file_path, each_file) submissions.append(print_job(copy)) # Add a submission without constraints submissions.append(print_job(deepcopy(fnew))) if get_yes_no_terminal("Will submit the previous entries. Press y/n to submit: ") == "y": for submission in submissions: print("Submitting job: %s" % submission['name']) submit_job(scanned_files=submission) else: print("Aborting submission.") def get_file(file_type): return PyZenity.GetFilename(title="Please select the %s file." % file_type, timeout=0, width=400)[0] def get_yes_no_terminal(message): response = input(message) if response != "y" and response != "n": print("Please enter 'y' or 'n'.") return get_yes_no_terminal(message) return response def get_yes_no(message, title="Please make a selection."): return PyZenity.Question(text=message, title=title, timeout=0, width=400) def get_text(message, default, title="Please make a selection.", choices=None): response = PyZenity.GetText(text=message, entry_text=default, title=title, timeout=0, width=400) if response is None: sys.exit(0) while choices and response not in choices: response = PyZenity.GetText(title="You didn't enter a valid value.", text=message, entry_text=default, timeout=0, width=400) if response is None: sys.exit(0) return response def monitor_job(key): while True: r = requests.get(options.server_url + "/entry/%s?return_type=json" % key) # If there is an error, raise it r.raise_for_status() r = r.json() if r['status'] == "complete": print("Entry complete! Opening in browser.") webbrowser.open(r['url']) sys.exit(0) else: print(r) time.sleep(30) # The possible return values are # {"status":"queued"} # {"status":"running", "stage":"fragment generation"} # {"status":"running", "stage":"starting abinitio"} # {"status":"running", "stage":"abinitio", "percent_complete":"xx.xx%"} # {"status":"running", "stage":"cleanup"} # {"status":"failure", "error": error} # {"status":"complete","url":"url_to_view_results"} def submit_job(scanned_files=None): if options.interactive: print("An interactive selection box should have appeared. If you don't see it, please make sure you don't " "have any windows set to 'Always on top'.") PyZenity.InfoMessage(title="Notice", text="Only academic submissions are allowed on this server.", width=400) ftype = get_text("Enter chemical shift file type: star/talos", "talos", choices=["star", "talos"]) if ftype == "star": options.star_file = get_file("chemical shift") elif ftype == "talos": options.talos_file = get_file("chemical shift") options.auth_email = get_text("Please enter your e-mail: ", "") options.prot_name = get_text("Please enter the protein name: ", "") options.auth_name = get_text("Please enter your name: ", getpass.getuser()) options.trim_type = get_text('Enter trim type: lange/notrim', "lange", choices=["lange", "notrim"]) options.structures = get_text("Please enter the number of structures to generate: ", "30000") if get_yes_no("Do you have disulfide bond, constraint, or RDC files to upload?"): if get_yes_no("Do you have a disulfide bond file?"): options.disulfide_file = get_file("Disulfide bond file:") if get_yes_no("Do you have an RDC file?"): options.rdc_file = get_file("RDC file:") if get_yes_no("Do you have a constraint file?"): options.constraint_file = get_file("Constraint file:") message = """ Here is what you have entered: Author name: %s Author email: %s Protein name: %s Num structures: %s Trim type: %s Chemical shift type: %s Chemical shift file: %s Constraint file: %s Disulfide bond file: %s RDC file: %s Is this correct?""" % (options.auth_name, options.auth_email, options.prot_name, options.structures, options.trim_type, "star" if options.star_file else "talos", options.star_file if options.star_file else options.talos_file, options.constraint_file, options.disulfide_file, options.rdc_file) if not get_yes_no(message): PyZenity.ErrorMessage(text="Exiting because you said something was wrong.", width=400) sys.exit(0) # Structures must be specified in increments of 25 (and no more than 40000). # Generate the payload depending on whether we are provided files already payload = { 'auth_email': options.auth_email, 'structures': options.structures, 'prot_name': options.prot_name, 'auth_name': options.auth_name, 'trim_type': options.trim_type, 'return_type': 'json', 'academic_checkbox': 'checked' } files = {} # Folder scanning submission if scanned_files: if 'star' in scanned_files: payload['chemical_shift_type'] = 'star' files['cs_file'] = open(scanned_files['star'], 'rb') elif 'talos' in scanned_files: payload['chemical_shift_type'] = 'talos' # Upload the chemical shift file (in TALOS format) files['cs_file'] = open(scanned_files['talos'], 'rb') else: raise ValueError("You must provide either 'star' or 'talos' file in the scanned_files.") if 'constraints_file' in scanned_files: files['constraints_file'] = open(scanned_files['constraints_file'], 'rb') if 'rdc_file' in scanned_files: files['rdc_file'] = open(scanned_files['rdc_file'], 'rb') if 'disulfide_file' in scanned_files: files['disulfide_file'] = open(scanned_files['disulfide_file'], 'rb') if 'name' in scanned_files: payload['prot_name'] = scanned_files['name'] # "Normal" submission else: if options.star_file: payload['chemical_shift_type'] = 'star' # Upload the chemical shift file (in STAR format) files['cs_file'] = open(options.star_file, 'rb') elif options.talos_file: payload['chemical_shift_type'] = 'talos' # Upload the chemical shift file (in TALOS format) files['cs_file'] = open(options.talos_file, 'rb') else: print("You must specify either a STAR or TALOS file.") sys.exit(4) # Optionally, include disulfide bond, risidual dipolar coupling, or Rosetta Constraint file if options.rdc_file: files['rdc_file'] = open(options.rdc_file, 'rb') if options.constraint_file: files['constraints_file'] = open(options.constraint_file, 'rb') if options.disulfide_file: files['disulfide_file'] = open(options.disulfide_file, 'rb') # Submit the request r = requests.post(options.server_url, data=payload, files=files) # If there is an HTTP error code, raise it r.raise_for_status() # You are returned an error or success code as json. If success, you will also be returned the entry key result = r.json() # Check the status if result['status'] != "success": if options.interactive: PyZenity.ErrorMessage(text="At least one error occurred! Errors:\n" + "\n".join(result['errors']), width=400) print("At least one error occurred! Errors:\n" + "\n".join(result['errors'])) return 'err' # Successfully submission else: if options.interactive: try: import webbrowser webbrowser.open(options.server_url + "/entry/%s" % result['key']) except ImportError: pass PyZenity.InfoMessage(text="Entry submitted with key: %s\n" % result['key'], width=400) print("Entry submitted with key: %s\nGo to: https://csrosetta.bmrb.io/entry/%s" % ( result['key'], result['key'])) return result['key'] # Monitor a job if asked if options.key is not None: monitor_job(options.key) sys.exit(0) # Submit a new job if len(args) != 0: print("Please use the options.") sys.exit(3) # If they want us to auto-submit directories then do that if options.directory_mode: scan_folder_for_jobs(options.directory_mode) sys.exit(0) # Submit a new job if submit_job() == 'err': sys.exit(1) else: sys.exit(0)