## hatch_ci.code

src/hatch_ci/code.py
    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