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
|