xref: /aosp_15_r20/external/pigweed/pw_env_setup/py/pw_env_setup/virtualenv_setup/install.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2020 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://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 "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Sets up a Python 3 virtualenv for Pigweed."""
15
16import contextlib
17import datetime
18import glob
19import os
20import platform
21import re
22import shutil
23import subprocess
24import sys
25import stat
26import tempfile
27
28# Grabbing datetime string once so it will always be the same for all GnTarget
29# objects.
30_DATETIME_STRING = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
31
32
33def _is_windows() -> bool:
34    return platform.system().lower() == 'windows'
35
36
37class GnTarget:
38    def __init__(self, val):
39        self.directory, self.target = val.split('#', 1)
40        self.name = '-'.join(
41            (re.sub(r'\W+', '_', self.target).strip('_'), _DATETIME_STRING)
42        )
43
44
45def git_stdout(*args, **kwargs):
46    """Run git, passing args as git params and kwargs to subprocess."""
47    return subprocess.check_output(['git'] + list(args), **kwargs).strip()
48
49
50def git_repo_root(path='./'):
51    """Find git repository root."""
52    try:
53        return git_stdout('-C', path, 'rev-parse', '--show-toplevel')
54    except subprocess.CalledProcessError:
55        return None
56
57
58class GitRepoNotFound(Exception):
59    """Git repository not found."""
60
61
62def _installed_packages(venv_python):
63    cmd = (venv_python, '-m', 'pip', '--disable-pip-version-check', 'list')
64    output = subprocess.check_output(cmd).splitlines()
65    return set(x.split()[0].lower() for x in output[2:])
66
67
68def _required_packages(requirements):
69    packages = set()
70
71    for req in requirements:
72        with open(req, 'r') as ins:
73            for line in ins:
74                line = line.strip()
75                if not line or line.startswith('#'):
76                    continue
77                packages.add(line.split('=')[0])
78
79    return packages
80
81
82def _check_call(args, **kwargs):
83    stdout = kwargs.get('stdout', sys.stdout)
84
85    with tempfile.TemporaryFile(mode='w+') as temp:
86        try:
87            kwargs['stdout'] = temp
88            kwargs['stderr'] = subprocess.STDOUT
89            print(args, kwargs, file=temp)
90            subprocess.check_call(args, **kwargs)
91        except subprocess.CalledProcessError:
92            temp.seek(0)
93            stdout.write(temp.read())
94            raise
95
96
97def _find_files_by_name(roots, name, allow_nesting=False):
98    matches = []
99    for root in roots:
100        for dirpart, dirs, files in os.walk(root):
101            if name in files:
102                matches.append(os.path.join(dirpart, name))
103                # If this directory is a match don't recurse inside it looking
104                # for more matches.
105                if not allow_nesting:
106                    dirs[:] = []
107
108            # Filter directories starting with . to avoid searching unnecessary
109            # paths and finding files that should be hidden.
110            dirs[:] = [d for d in dirs if not d.startswith('.')]
111    return matches
112
113
114def _check_venv(python, version, venv_path, pyvenv_cfg):
115    if _is_windows():
116        return
117
118    # Check if the python location and version used for the existing virtualenv
119    # is the same as the python we're using. If it doesn't match, we need to
120    # delete the existing virtualenv and start again.
121    if os.path.exists(pyvenv_cfg):
122        pyvenv_values = {}
123        with open(pyvenv_cfg, 'r') as ins:
124            for line in ins:
125                key, value = line.strip().split(' = ', 1)
126                pyvenv_values[key] = value
127        pydir = os.path.dirname(python)
128        home = pyvenv_values.get('home')
129        if pydir != home and not pydir.startswith(venv_path):
130            shutil.rmtree(venv_path)
131        elif pyvenv_values.get('version') not in '.'.join(map(str, version)):
132            shutil.rmtree(venv_path)
133
134
135def _check_python_install_permissions(python):
136    # These pickle files are not included on windows.
137    # The path on windows is environment/cipd/packages/python/bin/Lib/lib2to3/
138    if _is_windows():
139        return
140
141    # Make any existing lib2to3 pickle files read+write. This is needed for
142    # importing yapf.
143    lib2to3_path = os.path.join(
144        os.path.dirname(os.path.dirname(python)), 'lib', 'python3.9', 'lib2to3'
145    )
146    pickle_file_paths = []
147    if os.path.isdir(lib2to3_path):
148        pickle_file_paths.extend(
149            file_path
150            for file_path in os.listdir(lib2to3_path)
151            if '.pickle' in file_path
152        )
153    try:
154        for pickle_file in pickle_file_paths:
155            pickle_full_path = os.path.join(lib2to3_path, pickle_file)
156            os.chmod(
157                pickle_full_path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP
158            )
159    except PermissionError:
160        pass
161
162
163def _flatten(*items):
164    """Yields items from a series of items and nested iterables."""
165
166    for item in items:
167        if isinstance(item, (list, tuple)):
168            for i in _flatten(*item):
169                yield i
170        else:
171            yield item
172
173
174def _python_version(python_path: str):
175    """Returns the version (major, minor, rev) of the `python_path` binary."""
176    # Prints values like "3.10.0"
177    command = (
178        python_path,
179        '-c',
180        'import sys; print(".".join(map(str, sys.version_info[:3])))',
181    )
182    version_str = (
183        subprocess.check_output(command, stderr=subprocess.STDOUT)
184        .strip()
185        .decode()
186    )
187    return tuple(map(int, version_str.split('.')))
188
189
190def install(  # pylint: disable=too-many-arguments,too-many-locals,too-many-statements
191    project_root,
192    venv_path,
193    full_envsetup=True,
194    requirements=None,
195    constraints=None,
196    pip_install_disable_cache=None,
197    pip_install_find_links=None,
198    pip_install_offline=None,
199    pip_install_require_hashes=None,
200    gn_args=(),
201    gn_targets=(),
202    gn_out_dir=None,
203    python=sys.executable,
204    env=None,
205    system_packages=False,
206    use_pinned_pip_packages=True,
207):
208    """Creates a venv and installs all packages in this Git repo."""
209
210    version = _python_version(python)
211    if version[0] != 3:
212        print('=' * 60, file=sys.stderr)
213        print('Unexpected Python version:', version, file=sys.stderr)
214        print('=' * 60, file=sys.stderr)
215        return False
216
217    # The bin/ directory is called Scripts/ on Windows. Don't ask.
218    venv_bin = os.path.join(venv_path, 'Scripts' if os.name == 'nt' else 'bin')
219
220    if env:
221        env.set('VIRTUAL_ENV', venv_path)
222        env.prepend('PATH', venv_bin)
223        env.clear('PYTHONHOME')
224        env.clear('__PYVENV_LAUNCHER__')
225    else:
226        env = contextlib.nullcontext()
227
228    # Virtual environments may contain read-only files (notably activate
229    # scripts).  `venv` calls below will fail if they are not writeable.
230    if os.path.isdir(venv_path):
231        for root, _dirs, files in os.walk(venv_path):
232            for file in files:
233                path = os.path.join(root, file)
234                mode = os.lstat(path).st_mode
235                if not (stat.S_ISLNK(mode) or (mode & stat.S_IWRITE)):
236                    os.chmod(path, mode | stat.S_IWRITE)
237
238    pyvenv_cfg = os.path.join(venv_path, 'pyvenv.cfg')
239
240    _check_python_install_permissions(python)
241    _check_venv(python, version, venv_path, pyvenv_cfg)
242
243    if full_envsetup or not os.path.exists(pyvenv_cfg):
244        # On Mac sometimes the CIPD Python has __PYVENV_LAUNCHER__ set to
245        # point to the system Python, which causes CIPD Python to create
246        # virtualenvs that reference the system Python instead of the CIPD
247        # Python. Clearing __PYVENV_LAUNCHER__ fixes that. See also pwbug/59.
248        envcopy = os.environ.copy()
249        if '__PYVENV_LAUNCHER__' in envcopy:
250            del envcopy['__PYVENV_LAUNCHER__']
251
252        # TODO(spang): Pass --upgrade-deps and remove pip & setuptools
253        # upgrade below. This can only be done once the minimum python
254        # version is at least 3.9.
255        cmd = [python, '-m', 'venv']
256
257        # Windows requires strange wizardry, and must follow symlinks
258        # starting with 3.11.
259        #
260        # Without this, windows fails bootstrap trying to copy
261        # "environment\cipd\packages\python\bin\venvlauncher.exe"
262        #
263        # This file doesn't exist in Python 3.11 on Windows and may be a bug
264        # in venv. Pigweed already uses symlinks on Windows for the GN build,
265        # so adding this option is not an issue.
266        #
267        # Further excitement is had when trying to update a virtual environment
268        # that is created using symlinks under Windows.  `venv` will fail with
269        # and error that the source and destination are the same file.  To work
270        # around this, we run `venv` in `--clear` mode under Windows.
271        if _is_windows() and version >= (3, 11):
272            cmd += ['--clear', '--symlinks']
273        else:
274            cmd += ['--upgrade']
275
276        cmd += ['--system-site-packages'] if system_packages else []
277        cmd += [venv_path]
278        _check_call(cmd, env=envcopy)
279
280    venv_python = os.path.join(venv_bin, 'python')
281
282    pw_root = os.environ.get('PW_ROOT')
283    if not pw_root and env:
284        pw_root = env.PW_ROOT
285    if not pw_root:
286        pw_root = git_repo_root()
287    if not pw_root:
288        raise GitRepoNotFound()
289
290    # Sometimes we get an error saying "Egg-link ... does not match
291    # installed location". This gets around that. The egg-link files
292    # all come from 'pw'-prefixed packages we installed with --editable.
293    # Source: https://stackoverflow.com/a/48972085
294    for egg_link in glob.glob(
295        os.path.join(venv_path, 'lib/python*/site-packages/*.egg-link')
296    ):
297        os.unlink(egg_link)
298
299    pip_install_args = []
300    if pip_install_find_links:
301        for package_dir in pip_install_find_links:
302            pip_install_args.append('--find-links')
303            with env():
304                pip_install_args.append(os.path.expandvars(package_dir))
305    if pip_install_require_hashes:
306        pip_install_args.append('--require-hashes')
307    if pip_install_offline:
308        pip_install_args.append('--no-index')
309    if pip_install_disable_cache:
310        pip_install_args.append('--no-cache-dir')
311
312    def pip_install(*args):
313        args = list(_flatten(args))
314        with env():
315            cmd = (
316                [
317                    venv_python,
318                    '-m',
319                    'pip',
320                    '--disable-pip-version-check',
321                    'install',
322                ]
323                + pip_install_args
324                + args
325            )
326            return _check_call(cmd)
327
328    constraint_args = []
329    if constraints:
330        constraint_args.extend(
331            '--constraint={}'.format(constraint) for constraint in constraints
332        )
333
334    pip_install(
335        '--log',
336        os.path.join(venv_path, 'pip-upgrade.log'),
337        '--upgrade',
338        'pip',
339        'setuptools',
340        'toml',  # Needed for pyproject.toml package installs.
341        # Include wheel so pip installs can be done without build
342        # isolation.
343        'wheel',
344        'pip-tools',
345        constraint_args,
346    )
347
348    # TODO(tonymd): Remove this when projects have defined requirements.
349    if (not requirements) and constraints:
350        requirements = constraints
351
352    if requirements:
353        requirement_args = []
354        # Note: --no-build-isolation should be avoided for installing 3rd party
355        # Python packages that use C/C++ extension modules.
356        # https://setuptools.pypa.io/en/latest/userguide/ext_modules.html
357        requirement_args.extend(
358            '--requirement={}'.format(req) for req in requirements
359        )
360        combined_requirement_args = requirement_args + constraint_args
361        pip_install(
362            '--log',
363            os.path.join(venv_path, 'pip-requirements.log'),
364            combined_requirement_args,
365        )
366
367    def install_packages(gn_target):
368        if gn_out_dir is None:
369            build_dir = os.path.join(venv_path, 'gn')
370        else:
371            build_dir = gn_out_dir
372
373        env_log = 'env-{}.log'.format(gn_target.name)
374        env_log_path = os.path.join(venv_path, env_log)
375        with open(env_log_path, 'w') as outs:
376            for key, value in sorted(os.environ.items()):
377                if key.upper().endswith('PATH'):
378                    print(key, '=', file=outs)
379                    # pylint: disable=invalid-name
380                    for v in value.split(os.pathsep):
381                        print('   ', v, file=outs)
382                    # pylint: enable=invalid-name
383                else:
384                    print(key, '=', value, file=outs)
385
386        gn_log = 'gn-gen-{}.log'.format(gn_target.name)
387        gn_log_path = os.path.join(venv_path, gn_log)
388        try:
389            with open(gn_log_path, 'w') as outs:
390                gn_cmd = ['gn', 'gen', build_dir]
391
392                args = list(gn_args)
393                if not use_pinned_pip_packages:
394                    args.append('pw_build_PIP_CONSTRAINTS=[]')
395
396                args.append('dir_pigweed="{}"'.format(pw_root))
397                gn_cmd.append('--args={}'.format(' '.join(args)))
398
399                print(gn_cmd, file=outs)
400                subprocess.check_call(
401                    gn_cmd,
402                    cwd=os.path.join(project_root, gn_target.directory),
403                    stdout=outs,
404                    stderr=outs,
405                )
406        except subprocess.CalledProcessError as err:
407            with open(gn_log_path, 'r') as ins:
408                raise subprocess.CalledProcessError(
409                    err.returncode, err.cmd, ins.read()
410                )
411
412        ninja_log = 'ninja-{}.log'.format(gn_target.name)
413        ninja_log_path = os.path.join(venv_path, ninja_log)
414        try:
415            with open(ninja_log_path, 'w') as outs:
416                ninja_cmd = ['ninja', '-C', build_dir, '-v']
417                ninja_cmd.append(gn_target.target)
418                print(ninja_cmd, file=outs)
419                subprocess.check_call(ninja_cmd, stdout=outs, stderr=outs)
420        except subprocess.CalledProcessError as err:
421            with open(ninja_log_path, 'r') as ins:
422                raise subprocess.CalledProcessError(
423                    err.returncode, err.cmd, ins.read()
424                )
425
426        with open(os.path.join(venv_path, 'pip-list.log'), 'w') as outs:
427            subprocess.check_call(
428                [
429                    venv_python,
430                    '-m',
431                    'pip',
432                    '--disable-pip-version-check',
433                    'list',
434                ],
435                stdout=outs,
436            )
437
438    if gn_targets:
439        with env():
440            for gn_target in gn_targets:
441                install_packages(gn_target)
442
443    return True
444