File: Check.py

package info (click to toggle)
python-azure 20251104%2Bgit-1
  • links: PTS, VCS
  • area: main
  • in suites: forky
  • size: 770,224 kB
  • sloc: python: 6,357,217; ansic: 804; javascript: 287; makefile: 198; sh: 193; xml: 109
file content (244 lines) | stat: -rw-r--r-- 10,344 bytes parent folder | download | duplicates (2)
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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
import abc
import os
import argparse
import traceback
import sys
import shutil
import tempfile
import pathlib

from typing import Sequence, Optional, List, Any, Tuple
import subprocess

from ci_tools.parsing import ParsedSetup
from ci_tools.functions import (
    discover_targeted_packages,
    get_venv_call,
    install_into_venv,
    get_venv_python,
    get_pip_command,
    find_whl,
)
from ci_tools.variables import discover_repo_root, in_ci
from ci_tools.logging import logger

# right now, we are assuming you HAVE to be in the azure-sdk-tools repo
# we assume this because we don't know how a dev has installed this package, and might be
# being called from within a site-packages folder. Due to that, we can't trust the location of __file__
REPO_ROOT = discover_repo_root()


class Check(abc.ABC):
    """
    Base class for checks.

    Subclasses must implement register() to add a subparser for the check.
    """

    def __init__(self) -> None:
        pass

    @abc.abstractmethod
    def register(
        self, subparsers: "argparse._SubParsersAction", parent_parsers: Optional[List[argparse.ArgumentParser]] = None
    ) -> None:
        """
        Register this check with the CLI subparsers.

        `subparsers` is the object returned by ArgumentParser.add_subparsers().
        `parent_parsers` can be a list of ArgumentParser objects to be used as parents.
        Subclasses MUST implement this method.
        """
        raise NotImplementedError

    def run(self, args: argparse.Namespace) -> int:
        """Run the check command.

        Subclasses can override this to perform the actual work.
        """
        return 0

    def create_venv(self, isolate: bool, venv_location: str) -> str:
        """Abstraction for creating a virtual environment."""
        if isolate:
            venv_cmd = get_venv_call(sys.executable)
            venv_python = get_venv_python(venv_location)
            if os.path.exists(venv_python):
                logger.info(f"Reusing existing venv at {venv_python}")
                return venv_python
            else:
                shutil.rmtree(venv_location, ignore_errors=True)

            subprocess.check_call(venv_cmd + [venv_location])

            if in_ci():
                # first attempt to retrieve azure-sdk-tools from the prebuilt wheel directory
                # if present, install from there instead of constantly rebuilding azure-sdk-tools in a possible
                # parallel situation
                wheel_dir = os.getenv("PREBUILT_WHEEL_DIR", None) or os.path.join(REPO_ROOT, ".wheels")
                prebuilt_whl = find_whl(wheel_dir, "azure-sdk-tools", "0.0.0")

                if prebuilt_whl:
                    install_location = os.path.join(wheel_dir, prebuilt_whl)
                    install_into_venv(venv_location, [f"{install_location}[build]"], REPO_ROOT)
                else:
                    logger.error(
                        "Falling back to manual build and install of azure-sdk-tools into isolated env,"
                        f" unable to locate prebuilt azure-sdk-tools within {wheel_dir}"
                    )
            else:
                install_into_venv(
                    venv_location, [os.path.join(REPO_ROOT, "eng/tools/azure-sdk-tools[build]")], REPO_ROOT
                )

            venv_python_exe = get_venv_python(venv_location)

            return venv_python_exe

        # if we don't need to isolate, just return the python executable that we're invoking
        return sys.executable

    def get_executable(self, isolate: bool, check_name: str, executable: str, package_folder: str) -> Tuple[str, str]:
        """Get the Python executable that should be used for this check."""
        venv_location = os.path.join(package_folder, f".venv_{check_name}")

        # if isolation is required, the executable we get back will align with the venv
        # otherwise we'll just get sys.executable and install in current
        executable = self.create_venv(isolate, venv_location)
        staging_directory = os.path.join(venv_location, ".staging")
        os.makedirs(staging_directory, exist_ok=True)
        return executable, staging_directory

    def run_venv_command(
        self, executable: str, command: Sequence[str], cwd: str, check: bool = False, append_executable: bool = True
    ) -> subprocess.CompletedProcess[str]:
        """Run a command in the given virtual environment.
        - Prepends the virtual environment's bin directory to the PATH environment variable (if one exists)
        - Uses the provided Python executable to run the command.
        - Collects the output.
        - If check is True, raise CalledProcessError on failure."""

        if command[0].endswith("python") or command[0].endswith("python.exe"):
            raise ValueError(
                "The command array should not include the python executable, it is provided by the 'executable' argument"
            )

        env = os.environ.copy()

        python_exec = pathlib.Path(executable)
        if python_exec.exists():
            venv_bin = str(python_exec.parent)
            venv_root = str(python_exec.parent.parent)
            env["VIRTUAL_ENV"] = venv_root
            env["PATH"] = venv_bin + os.pathsep + env.get("PATH", "")
            env.pop("PYTHONPATH", None)
            env.pop("PYTHONHOME", None)
        else:
            raise RuntimeError(f"Unable to find parent venv for executable {executable}")

        # When not appending executable, resolve the command using the modified PATH
        if not append_executable:
            resolved = shutil.which(command[0], path=env["PATH"])
            if not resolved:
                raise RuntimeError(f"Command '{command[0]}' not found in PATH: {env['PATH']}")
            cmd_to_run = [resolved] + list(command[1:])
        else:
            cmd_to_run = [executable] + list(command)

        logger.debug(f"Running command: {cmd_to_run}.")
        logger.debug(f"VIRTUAL_ENV: {env['VIRTUAL_ENV']}.")
        logger.debug(f"PATH : {env['PATH']}.")

        result = subprocess.run(
            cmd_to_run, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=check, env=env
        )

        return result

    def get_targeted_directories(self, args: argparse.Namespace) -> List[ParsedSetup]:
        """
        Get the directories that are targeted for the check.
        """
        targeted: List[ParsedSetup] = []
        targeted_dir = os.getcwd()

        if args.target == ".":
            try:
                targeted.append(ParsedSetup.from_path(targeted_dir))
            except Exception as e:
                logger.error(
                    "Error: Current directory does not appear to be a Python package (no setup.py or setup.cfg found). Remove '.' argument to run on child directories."
                )
                logger.error(f"Exception: {e}")
                return []
        else:
            targeted_packages = discover_targeted_packages(args.target, targeted_dir)
            for pkg in targeted_packages:
                try:
                    targeted.append(ParsedSetup.from_path(pkg))
                except Exception as e:
                    logger.error(f"Unable to parse {pkg} as a Python package. Dumping exception detail and skipping.")
                    logger.error(f"Exception: {e}")
                    logger.error(traceback.format_exc())

        return targeted

    def install_dev_reqs(self, executable: str, args: argparse.Namespace, package_dir: str) -> None:
        """Install dev requirements for the given package."""
        dev_requirements = os.path.join(package_dir, "dev_requirements.txt")

        requirements = []
        if os.path.exists(dev_requirements):
            requirements += ["-r", dev_requirements]
        else:
            logger.warning(
                f"No dev_requirements.txt found for {package_dir}, skipping installation of dev requirements."
            )
            return

        temp_req_file = None
        if not getattr(args, "isolate", False):
            # don't install azure-sdk-tools when not isolated
            with open(dev_requirements, "r") as f:
                filtered_req_lines = [line.strip() for line in f if "eng/tools/azure-sdk-tools" not in line]
            with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_req_file:
                temp_req_file.write("\n".join(filtered_req_lines))
            if temp_req_file.name:
                requirements = ["-r", temp_req_file.name]
        try:
            logger.info(f"Installing dev requirements for {package_dir}")
            install_into_venv(executable, requirements, package_dir)
        except subprocess.CalledProcessError as e:
            logger.error("Failed to install dev requirements:", e)
            raise e
        finally:
            if temp_req_file and temp_req_file.name:
                try:
                    os.remove(temp_req_file.name)
                except Exception as cleanup_error:
                    logger.warning(f"Failed to remove temporary requirements file: {cleanup_error}")

    def pip_freeze(self, executable: str) -> None:
        """Run pip freeze in the given virtual environment and log the output. This function handles both isolated and non-isolated
        environments, as well as calling the proper `uv` executable with additional --python argument if needed.

        :param executable: Path to the python executable that should invoke this check.
        :returns None:
        """
        try:
            # to uv pip install or freeze to a target environment, we have to add `--python <path to python exe>`
            # to tell uv which environment to target
            command = get_pip_command(executable)

            if command[0] == "uv":
                command += ["freeze", "--python", executable]
            else:
                command += ["freeze"]

            result = subprocess.run(command, cwd=os.getcwd(), check=True, capture_output=True, text=True)
            logger.info("Installed packages:")
            logger.info(result.stdout)
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to run pip freeze: {e}")
            logger.error(e.stdout)
            logger.error(e.stderr)