xref: /aosp_15_r20/tools/netsim/scripts/environment.py (revision cf78ab8cffb8fc9207af348f23af247fb04370a6) 
1# Copyright 2024 - The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the',  help='License');
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an',  help='AS IS' BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14from collections import UserDict
15import json
16import logging
17import os
18from pathlib import Path
19import platform
20import subprocess
21from utils import rust_version
22
23
24class BaseEnvironment(UserDict):
25  """Base class for environment management, providing a common foundation for
26
27  both Posix and Windows environments.
28  """
29
30  def __init__(self, aosp: Path):
31    paths = [
32        str(
33            aosp
34            / "external"
35            / "qemu"
36            / "android"
37            / "third_party"
38            / "chromium"
39            / "depot_tools"
40        )
41    ]
42
43    # Append prebuilt rust toolchain except for darwin_aarch64
44    # TODO(360874898): aarch64-apple-darwin prebuilt rust toolchain supported from 1.77.1
45    if not (platform.system() == "Darwin" and platform.machine() == "arm64"):
46      paths.append(
47          str(
48              aosp
49              / "prebuilts"
50              / "rust"
51              / f"{platform.system().lower()}-x86"
52              / f"{rust_version()}"
53              / "bin"
54          )
55      )
56
57    paths.append(os.environ.get("PATH", ""))
58
59    super().__init__({"PATH": os.pathsep.join(paths)})
60
61
62class PosixEnvironment(BaseEnvironment):
63
64  def __init__(self, aosp: Path):
65    super().__init__(aosp)
66
67
68class VisualStudioNotFoundException(Exception):
69  pass
70
71
72class VisualStudioMissingVarException(Exception):
73  pass
74
75
76class VisualStudioNativeWorkloadNotFoundException(Exception):
77  pass
78
79
80class WindowsEnvironment(BaseEnvironment):
81  """Environment manager for Windows systems, specifically handling Visual Studio integration."""
82
83  def __init__(self, aosp: Path):
84    assert platform.system() == "Windows"
85    super().__init__(aosp)
86    for key in os.environ:
87      self[key.upper()] = os.environ[key]
88
89    vs = self._visual_studio()
90    logging.info("Loading environment from %s", vs)
91    env_lines = subprocess.check_output(
92        [vs, "&&", "set"], encoding="utf-8"
93    ).splitlines()
94    for line in env_lines:
95      if "=" in line:
96        key, val = line.split("=", 1)
97        # Variables in windows are case insensitive, but not in python dict!
98        self[key.upper()] = val
99
100    # Set PYTHONUTF8 to 1
101    self["PYTHONUTF8"] = "1"
102
103    if not "VSINSTALLDIR" in self:
104      raise VisualStudioMissingVarException(
105          "Missing VSINSTALLDIR in environment"
106      )
107
108    if not "VCTOOLSINSTALLDIR" in self:
109      raise VisualStudioMissingVarException(
110          "Missing VCTOOLSINSTALLDIR in environment"
111      )
112
113  def _visual_studio(self) -> Path:
114    """Locates the Visual Studio installation and its Native Desktop workload.
115
116    Raises:
117        VisualStudioNotFoundException: When Visual Studio is not found.
118        VisualStudioNativeWorkloadNotFoundException: When the Native Desktop
119        workload is not found.
120
121    Returns:
122        Path: Path to the Visual Studio vcvars64.bat file.
123    """
124    prgrfiles = Path(os.getenv("ProgramFiles(x86)", "C:\Program Files (x86)"))
125    res = subprocess.check_output([
126        str(
127            prgrfiles / "Microsoft Visual Studio" / "Installer" / "vswhere.exe"
128        ),
129        "-requires",
130        "Microsoft.VisualStudio.Workload.NativeDesktop",
131        "-sort",
132        "-format",
133        "json",
134        "-utf8",
135    ])
136    vsresult = json.loads(res)
137    if len(vsresult) == 0:
138      raise VisualStudioNativeWorkloadNotFoundException(
139          "No visual studio with the native desktop load available."
140      )
141
142    for install in vsresult:
143      logging.debug("Considering %s", install["displayName"])
144      candidates = list(
145          Path(install["installationPath"]).glob("**/vcvars64.bat")
146      )
147
148      if len(candidates) > 0:
149        return candidates[0].absolute()
150
151    # Oh oh, no visual studio..
152    raise VisualStudioNotFoundException(
153        "Unable to detect a visual studio installation with the native desktop"
154        " workload."
155    )
156
157
158def get_default_environment(aosp: Path):
159  """Returns the appropriate environment manager based on the current operating system.
160
161  The environment will make sure the following things hold:
162
163  - Ninja will be on the PATH
164  - The visual studio tools environment will be loaded
165  """
166  if platform.system() == "Windows":
167    return WindowsEnvironment(aosp)
168  return PosixEnvironment(aosp)
169