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
|
from __future__ import annotations
import ast
import re
from pathlib import Path
from typing import Any
class ValidationError(Exception):
pass
class MissingVariableError(ValidationError):
pass
def get_module_var(
path: Path | str, var: str = "__version__", abort=True
) -> str | None:
"""extract from a python module in path the module level <var> variable
Args:
path (str,Path): python module file to parse using ast (no code-execution)
var (str): module level variable name to extract
abort (bool): raise MissingVariable if var is not present
Returns:
None or str: the variable value if found or None
Raises:
MissingVariable: if the var is not found and abort is True
Notes:
this uses ast to parse path, so it doesn't load the module
"""
class V(ast.NodeVisitor):
def __init__(self, keys):
self.keys = keys
self.result = {}
def visit_Module(self, node): # noqa: N802
# we extract the module level variables
for subnode in ast.iter_child_nodes(node):
if not isinstance(subnode, ast.Assign):
continue
for target in subnode.targets:
if target.id not in self.keys:
continue
if not isinstance(subnode.value, (ast.Num, ast.Str, ast.Constant)):
raise ValidationError(
f"cannot extract non Constant variable "
f"{target.id} ({type(subnode.value)})"
)
if isinstance(subnode.value, ast.Str):
value = subnode.value.s
elif isinstance(subnode.value, ast.Num):
value = subnode.value.n
else:
value = subnode.value.value
if target.id in self.result:
raise ValidationError(
f"found multiple repeated variables {target.id}"
)
self.result[target.id] = value
return self.generic_visit(node)
v = V({var})
path = Path(path)
if path.exists():
tree = ast.parse(Path(path).read_text())
v.visit(tree)
if var not in v.result and abort:
raise MissingVariableError(f"cannot find {var} in {path}", path, var)
return v.result.get(var, None)
def set_module_var(
path: str | Path, var: str, value: Any, create: bool = True
) -> tuple[Any, str]:
"""replace var in path with value
Args:
path (str,Path): python module file to parse
var (str): module level variable name to extract
value (None or Any): if not None replace var in version_file
create (bool): create path if not present
Returns:
(str, str) the (<previous-var-value|None>, <the new text>)
"""
# validate the var
get_module_var(path, var, abort=False)
# module level var
expr = re.compile(f"^{var}\\s*=\\s*['\\\"](?P<value>[^\\\"']*)['\\\"]")
fixed = None
lines = []
src = Path(path)
if not src.exists() and create:
src.parent.mkdir(parents=True, exist_ok=True)
src.touch()
input_lines = src.read_text().split("\n")
for line in input_lines:
if fixed is not None:
lines.append(line)
continue
match = expr.search(line)
if match:
fixed = match.group("value")
if value is not None:
x, y = match.span(1)
line = line[:x] + value + line[y:]
lines.append(line)
txt = "\n".join(lines)
if (fixed is None) and create:
if txt and txt[-1] != "\n":
txt += "\n"
txt += f'{var} = "{value}"'
with Path(path).open("w") as fp:
fp.write(txt)
return fixed, txt
|