Coverage for /Users/jerry/Development/yenta/yenta/cli.py: 86%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

147 statements  

1#!/usr/bin/env python3 

2"""Console script for yenta.""" 

3import sys 

4import click 

5import configparser 

6import importlib.util 

7import more_itertools 

8import shutil 

9import os 

10 

11from rich.tree import Tree 

12from rich.text import Text 

13from rich import print 

14 

15from typing import Iterable 

16from networkx.drawing.nx_pydot import to_pydot 

17from colorama import init, Fore, Style 

18from pathlib import Path 

19from yenta.config import settings 

20from yenta.pipeline.Pipeline import Pipeline, TaskStatus 

21 

22import logging 

23 

24logger = logging.getLogger(__name__) 

25 

26 

27CHECK_MARK = u'\u2714' 

28X_MARK = u'\u2718' 

29 

30 

31def load_tasks(entry_file): 

32 spec = importlib.util.spec_from_file_location('main', entry_file) 

33 module = importlib.util.module_from_spec(spec) 

34 spec.loader.exec_module(module) 

35 

36 tasks = [func for _, func in module.__dict__.items() 

37 if callable(func) and hasattr(func, '_yenta_task')] 

38 

39 return tasks 

40 

41 

42@click.group() 

43@click.option('--config-file', default=settings.YENTA_CONFIG_FILE, type=Path, 

44 help='The config file from which to read settings.') 

45@click.option('--pipeline-store', type=Path, help='The directory to which the pipeline will be cached.') 

46@click.option('--entry-point', type=Path, help='The file containing the task definitions.') 

47@click.option('--log-file', type=Path, help='The file to which the logs should be written.') 

48def yenta(config_file, pipeline_store, entry_point, log_file): 

49 

50 init() 

51 

52 # append the local path we're running from so that we can allow 

53 # the project to import normally when running via CLI 

54 sys.path.append(os.getcwd()) 

55 

56 cf = configparser.ConfigParser() 

57 cf.read(config_file or settings.YENTA_CONFIG_FILE) 

58 if 'yenta' not in cf: 

59 cf['yenta'] = {} 

60 

61 settings.YENTA_ENTRY_POINT = entry_point or \ 

62 cf['yenta'].get('entry_point', None) or \ 

63 settings.YENTA_ENTRY_POINT 

64 

65 pipeline_file = cf['yenta'].get('pipeline_store', None) 

66 pipeline_path = Path(pipeline_file).resolve() if pipeline_file else None 

67 settings.YENTA_STORE_PATH = pipeline_store or \ 

68 pipeline_path or \ 

69 settings.YENTA_STORE_PATH 

70 conf_log_file = cf['yenta'].get('log_file', None) 

71 conf_log_path = Path(conf_log_file).resolve() if log_file else None 

72 settings.YENTA_LOG_FILE = log_file or \ 

73 conf_log_path or \ 

74 settings.YENTA_LOG_FILE 

75 

76 

77@yenta.command(help='List all available tasks.') 

78@click.option('--pipeline-name', default='default', help='The name of the pipeline to display.') 

79def list_tasks(pipeline_name='default'): 

80 

81 tasks = load_tasks(settings.YENTA_ENTRY_POINT) 

82 pipeline = Pipeline(*tasks) 

83 pipeline_data = Pipeline.load_pipeline(settings.YENTA_STORE_PATH / pipeline_name) 

84 

85 print('[bold white]The following tasks are available:[/bold white]') 

86 for task_name in pipeline.execution_order: 

87 task_result = pipeline_data.task_results.get(task_name, None) 

88 marker = ' ' 

89 if task_result and task_result.status == TaskStatus.SUCCESS: 

90 marker = f'[bold green]{CHECK_MARK}[/bold green]' 

91 elif task_result and task_result.status == TaskStatus.FAILURE: 

92 marker = f'[bold red]{X_MARK}[/bold red]' 

93 

94 print(f'[{marker}] [bold white]{task_name}[/bold white]') 

95 

96 

97@yenta.command(help='Show the current configuration.') 

98def show_config(): 

99 

100 print('[bold white]Yenta is using the following configuration:[/bold white]') 

101 print(f'[bold white]The entrypoint for Yenta is [green]{settings.YENTA_ENTRY_POINT}[/green][/bold white]') 

102 print(f'[bold white]Pipelines will be cached in [green]{settings.YENTA_STORE_PATH}[/green][/bold white]') 

103 if settings.YENTA_LOG_FILE: 

104 print('Log output will be written to ' + Fore.GREEN + str(settings.YENTA_LOG_FILE) + Fore.WHITE) 

105 else: 

106 print('No log output configured') 

107 

108 

109@yenta.command(help='Show information about a specific task.') 

110@click.argument('task-name') 

111@click.option('--pipeline-name', default='default', help='The name of the pipeline to display.') 

112def task_info(task_name, pipeline_name='default'): 

113 

114 tasks = load_tasks(settings.YENTA_ENTRY_POINT) 

115 pipeline_data = Pipeline.load_pipeline(settings.YENTA_STORE_PATH / pipeline_name) 

116 try: 

117 task = more_itertools.one(filter(lambda t: t.task_def.name == task_name, tasks)) 

118 print(f'[bold white]Information for task [green]{task_name}[/green]:[/bold white]') 

119 deps = ', '.join(task.task_def.depends_on) if task.task_def.depends_on else 'None' 

120 print(f'[bold white]Dependencies: [bright_blue]{deps}[/bright_blue][/bold white]') 

121 task_result = pipeline_data.task_results.get(task_name, None) 

122 marker = 'Did not run' 

123 if task_result and task_result.status == TaskStatus.SUCCESS: 

124 marker = Fore.GREEN + u'\u2714' + Fore.WHITE 

125 elif task_result and task_result.status == TaskStatus.FAILURE: 

126 marker = Fore.RED + u'\u2718' + ' ' + task_result.error + Fore.WHITE 

127 print(f'Previous status: {marker}') 

128 

129 tree = Tree('Previous result: ') 

130 if task_result and task_result.status == TaskStatus.SUCCESS: 

131 # print('Previous result: ', task_result) 

132 values_node = tree.add('values') 

133 for key in sorted(task_result.values.keys()): 

134 val = task_result.values.get(key) 

135 if isinstance(val, Iterable) and not isinstance(val, str): 

136 key_node = values_node.add(Text(f'{key}: ')) 

137 for v in val: 

138 key_node.add(Text(str(v))) 

139 else: 

140 values_node.add(Text(f'{key}: {val}')) 

141 artifacts_node = tree.add('artifacts') 

142 for key in sorted(task_result.artifacts.keys()): 

143 val = task_result.artifacts.get(key) 

144 if isinstance(val, Iterable) and not isinstance(val, str): 

145 key_node = artifacts_node.add(Text(f'{key}: ')) 

146 for v in val: 

147 key_node.add(Text(v)) 

148 else: 

149 artifacts_node.add(Text(f'{key}: {val}')) 

150 print(tree) 

151 else: 

152 print('Previous result: ' + Fore.GREEN + 'None' + Fore.WHITE) 

153 

154 except ValueError: 

155 print(Fore.WHITE + Style.BRIGHT + 'Unknown task ' + Fore.RED + task_name + Fore.WHITE + ' specified.') 

156 

157 

158@yenta.command(help='Remove a task from the pipeline cache.') 

159@click.argument('task-name') 

160@click.option('--pipeline-name', default='default', help='The name of the pipeline to display.') 

161def rm(task_name, pipeline_name='default'): 

162 

163 task_path = settings.YENTA_STORE_PATH / pipeline_name / task_name 

164 

165 if task_path.exists(): 

166 shutil.rmtree(task_path) 

167 else: 

168 print(Fore.WHITE + Style.BRIGHT + 'Unknown task ' + Fore.RED + task_name + Fore.WHITE + ' specified.') 

169 

170 

171@yenta.command(help='Mark a task as ignorable; it will be skipped by the pipeline.') 

172@click.argument('task-name') 

173@click.option('--pipeline-name', default='default', help='The name of the pipeline to display.') 

174def ignore(task_name, pipeline_name='default'): 

175 

176 tasks = load_tasks(settings.YENTA_ENTRY_POINT) 

177 target_task = [task for task in tasks if task.task_def.name == task_name] 

178 if len(target_task) == 0: 

179 print(f'[bold white]Unknown task {task_name} specified.[/bold white]') 

180 elif len(target_task) > 1: 

181 print(f'[bold red]Multiple tasks have the same name: {task_name}. Task names must be unique ' 

182 f'within a pipeline.[/bold red]') 

183 else: 

184 ignore_path = settings.YENTA_STORE_PATH / pipeline_name / task_name / '.ignore' 

185 ignore_path.touch(exist_ok=True) 

186 print(f'[bold white]Setting task {task_name} to be ignored.[/bold white]') 

187 

188 

189@yenta.command(help='Dump the task graph to a DOT file.') 

190@click.argument('filename', type=click.Path()) 

191def dump_task_graph(filename: Path): 

192 

193 tasks = load_tasks(settings.YENTA_ENTRY_POINT) 

194 pipeline = Pipeline(*tasks) 

195 pydot_graph = to_pydot(pipeline.task_graph) 

196 pydot_graph.write(filename) 

197 

198 

199@yenta.command(help='Run the pipeline.') 

200@click.option('--up-to', help='Optionally run the pipeline up to and including a given task.') 

201@click.option('--force-rerun', '-f', multiple=True, default=[], help='Force specified tasks to rerun.') 

202@click.option('--pipeline-name', default='default', help='The name of the pipeline to run.') 

203def run(up_to=None, force_rerun=None, pipeline_name='default'): 

204 

205 logger.info('Running the pipeline') 

206 tasks = load_tasks(settings.YENTA_ENTRY_POINT) 

207 pipeline = Pipeline(*tasks, name=pipeline_name) 

208 result = pipeline.run_pipeline(up_to, force_rerun) 

209 

210 

211if __name__ == "__main__": 

212 sys.exit(yenta()) # pragma: no cover