tools.py
3.88 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
'''
File: perguntations/tools.py
Description: Helper functions to load yaml files and run external programs.
'''
# python standard library
import asyncio
import logging
from os import path
import subprocess
from typing import Any, List
# third party libraries
import yaml
# setup logger for this module
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
def load_yaml(filename: str, default: Any = None) -> Any:
'''load data from yaml file'''
filename = path.expanduser(filename)
try:
file = open(filename, 'r', encoding='utf-8')
except Exception as exc:
logger.error(exc)
if default is not None:
return default
raise
with file:
try:
return yaml.safe_load(file)
except yaml.YAMLError as exc:
logger.error(str(exc).replace('\n', ' '))
if default is not None:
return default
raise
# ---------------------------------------------------------------------------
def run_script(script: str,
args: List[str],
stdin: str = '',
timeout: int = 3) -> Any:
'''
Runs a script and returns its stdout parsed as yaml, or None on error.
The script is run in another process but this function blocks waiting
for its termination.
'''
logger.info('run_script "%s"', script)
output = None
script = path.expanduser(script)
cmd = [script] + [str(a) for a in args]
# --- run process
try:
proc = subprocess.run(cmd,
input=stdin,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
universal_newlines=True,
timeout=timeout,
check=True,
)
except subprocess.TimeoutExpired:
logger.error('Timeout %ds exceeded running "%s".', timeout, script)
return output
except subprocess.CalledProcessError as exc:
logger.error('Return code %d running "%s".', exc.returncode, script)
return output
except OSError:
logger.error('Can not execute script "%s".', script)
return output
# --- parse yaml
try:
output = yaml.safe_load(proc.stdout)
except yaml.YAMLError:
logger.error('Error parsing yaml output of "%s".', script)
return output
# ----------------------------------------------------------------------------
async def run_script_async(script: str,
args: List[str],
stdin: str = '',
timeout: int = 3) -> Any:
'''Same as above, but asynchronous'''
script = path.expanduser(script)
args = [str(a) for a in args]
output = None
# --- start process
try:
proc = await asyncio.create_subprocess_exec(
script, *args,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.DEVNULL,
)
except OSError:
logger.error('Can not execute script "%s".', script)
return output
# --- send input and wait for termination
try:
stdout, _ = await asyncio.wait_for(
proc.communicate(input=stdin.encode('utf-8')),
timeout=timeout)
except asyncio.TimeoutError:
logger.warning('Timeout %ds running script "%s".', timeout, script)
return output
# --- check return code
if proc.returncode != 0:
logger.error('Return code %d running "%s".', proc.returncode, script)
return output
# --- parse yaml
try:
output = yaml.safe_load(stdout.decode('utf-8', 'ignore'))
except yaml.YAMLError:
logger.error('Error parsing yaml output of "%s"', script)
return output