tools.py
3.29 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
"""
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) -> Any:
"""load yaml file or raise exception on error"""
with open(path.expanduser(filename), "r", encoding="utf-8") as file:
return yaml.safe_load(file)
# ---------------------------------------------------------------------------
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.debug('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