1# coding: utf-8 2from __future__ import unicode_literals, division, absolute_import, print_function 3 4import ast 5import _ast 6import os 7import sys 8 9from . import package_root, task_keyword_args 10from ._import import _import_from 11 12 13if sys.version_info < (3,): 14 byte_cls = str 15else: 16 byte_cls = bytes 17 18 19def _list_tasks(): 20 """ 21 Fetches a list of all valid tasks that may be run, and the args they 22 accept. Does not actually import the task module to prevent errors if a 23 user does not have the dependencies installed for every task. 24 25 :return: 26 A list of 2-element tuples: 27 0: a unicode string of the task name 28 1: a list of dicts containing the parameter definitions 29 """ 30 31 out = [] 32 dev_path = os.path.join(package_root, 'dev') 33 for fname in sorted(os.listdir(dev_path)): 34 if fname.startswith('.') or fname.startswith('_'): 35 continue 36 if not fname.endswith('.py'): 37 continue 38 name = fname[:-3] 39 args = () 40 41 full_path = os.path.join(package_root, 'dev', fname) 42 with open(full_path, 'rb') as f: 43 full_code = f.read() 44 if sys.version_info >= (3,): 45 full_code = full_code.decode('utf-8') 46 47 task_node = ast.parse(full_code, filename=full_path) 48 for node in ast.iter_child_nodes(task_node): 49 if isinstance(node, _ast.Assign): 50 if len(node.targets) == 1 \ 51 and isinstance(node.targets[0], _ast.Name) \ 52 and node.targets[0].id == 'run_args': 53 args = ast.literal_eval(node.value) 54 break 55 56 out.append((name, args)) 57 return out 58 59 60def show_usage(): 61 """ 62 Prints to stderr the valid options for invoking tasks 63 """ 64 65 valid_tasks = [] 66 for task in _list_tasks(): 67 usage = task[0] 68 for run_arg in task[1]: 69 usage += ' ' 70 name = run_arg.get('name', '') 71 if run_arg.get('required', False): 72 usage += '{%s}' % name 73 else: 74 usage += '[%s]' % name 75 valid_tasks.append(usage) 76 77 out = 'Usage: run.py' 78 for karg in task_keyword_args: 79 out += ' [%s=%s]' % (karg['name'], karg['placeholder']) 80 out += ' (%s)' % ' | '.join(valid_tasks) 81 82 print(out, file=sys.stderr) 83 sys.exit(1) 84 85 86def _get_arg(num): 87 """ 88 :return: 89 A unicode string of the requested command line arg 90 """ 91 92 if len(sys.argv) < num + 1: 93 return None 94 arg = sys.argv[num] 95 if isinstance(arg, byte_cls): 96 arg = arg.decode('utf-8') 97 return arg 98 99 100def run_task(): 101 """ 102 Parses the command line args, invoking the requested task 103 """ 104 105 arg_num = 1 106 task = None 107 args = [] 108 kwargs = {} 109 110 # We look for the task name, processing any global task keyword args 111 # by setting the appropriate env var 112 while True: 113 val = _get_arg(arg_num) 114 if val is None: 115 break 116 117 next_arg = False 118 for karg in task_keyword_args: 119 if val.startswith(karg['name'] + '='): 120 os.environ[karg['env_var']] = val[len(karg['name']) + 1:] 121 next_arg = True 122 break 123 124 if next_arg: 125 arg_num += 1 126 continue 127 128 task = val 129 break 130 131 if task is None: 132 show_usage() 133 134 task_mod = _import_from('dev.%s' % task, package_root, allow_error=True) 135 if task_mod is None: 136 show_usage() 137 138 run_args = task_mod.__dict__.get('run_args', []) 139 max_args = arg_num + 1 + len(run_args) 140 141 if len(sys.argv) > max_args: 142 show_usage() 143 144 for i, run_arg in enumerate(run_args): 145 val = _get_arg(arg_num + 1 + i) 146 if val is None: 147 if run_arg.get('required', False): 148 show_usage() 149 break 150 151 if run_arg.get('cast') == 'int' and val.isdigit(): 152 val = int(val) 153 154 kwarg = run_arg.get('kwarg') 155 if kwarg: 156 kwargs[kwarg] = val 157 else: 158 args.append(val) 159 160 run = task_mod.__dict__.get('run') 161 162 result = run(*args, **kwargs) 163 sys.exit(int(not result)) 164