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
16from __future__ import print_function
17
18import glob
19import hashlib
20import os
21import re
22import subprocess
23import sys
24import tempfile
25
26
27class GnTarget(object):  # pylint: disable=useless-object-inheritance
28    def __init__(self, val):
29        self.directory, self.target = val.split('#', 1)
30        # hash() doesn't necessarily give the same value in new runs of Python,
31        # so compute a unique id for this object that's consistent from run to
32        # run.
33        try:
34            val = val.encode()
35        except AttributeError:
36            pass
37        self._unique_id = hashlib.md5(val).hexdigest()
38
39    @property
40    def name(self):
41        """A reasonably stable and unique name for each pair."""
42        result = '{}-{}'.format(
43            os.path.basename(os.path.normpath(self.directory)),
44            self._unique_id)
45        return re.sub(r'[:/#_]+', '_', result)
46
47
48def git_stdout(*args, **kwargs):
49    """Run git, passing args as git params and kwargs to subprocess."""
50    return subprocess.check_output(['git'] + list(args), **kwargs).strip()
51
52
53def git_repo_root(path='./'):
54    """Find git repository root."""
55    try:
56        return git_stdout('-C', path, 'rev-parse', '--show-toplevel')
57    except subprocess.CalledProcessError:
58        return None
59
60
61class GitRepoNotFound(Exception):
62    """Git repository not found."""
63
64
65def _installed_packages(venv_python):
66    cmd = (venv_python, '-m', 'pip', 'list', '--disable-pip-version-check')
67    output = subprocess.check_output(cmd).splitlines()
68    return set(x.split()[0].lower() for x in output[2:])
69
70
71def _required_packages(requirements):
72    packages = set()
73
74    for req in requirements:
75        with open(req, 'r') as ins:
76            for line in ins:
77                line = line.strip()
78                if not line or line.startswith('#'):
79                    continue
80                packages.add(line.split('=')[0])
81
82    return packages
83
84
85# TODO(pwbug/135) Move to common utility module.
86def _check_call(args, **kwargs):
87    stdout = kwargs.get('stdout', sys.stdout)
88
89    with tempfile.TemporaryFile(mode='w+') as temp:
90        try:
91            kwargs['stdout'] = temp
92            kwargs['stderr'] = subprocess.STDOUT
93            print(args, kwargs, file=temp)
94            subprocess.check_call(args, **kwargs)
95        except subprocess.CalledProcessError:
96            temp.seek(0)
97            stdout.write(temp.read())
98            raise
99
100
101def _find_files_by_name(roots, name, allow_nesting=False):
102    matches = []
103    for root in roots:
104        for dirpart, dirs, files in os.walk(root):
105            if name in files:
106                matches.append(os.path.join(dirpart, name))
107                # If this directory is a match don't recurse inside it looking
108                # for more matches.
109                if not allow_nesting:
110                    dirs[:] = []
111
112            # Filter directories starting with . to avoid searching unnecessary
113            # paths and finding files that should be hidden.
114            dirs[:] = [d for d in dirs if not d.startswith('.')]
115    return matches
116
117
118def install(
119        project_root,
120        venv_path,
121        full_envsetup=True,
122        requirements=(),
123        gn_targets=(),
124        gn_out_dir=None,
125        python=sys.executable,
126        env=None,
127):
128    """Creates a venv and installs all packages in this Git repo."""
129
130    version = subprocess.check_output(
131        (python, '--version'), stderr=subprocess.STDOUT).strip().decode()
132    # We expect Python 3.8, but if it came from CIPD let it pass anyway.
133    if ('3.8' not in version and '3.9' not in version
134            and 'chromium' not in version):
135        print('=' * 60, file=sys.stderr)
136        print('Unexpected Python version:', version, file=sys.stderr)
137        print('=' * 60, file=sys.stderr)
138        return False
139
140    # The bin/ directory is called Scripts/ on Windows. Don't ask.
141    venv_bin = os.path.join(venv_path, 'Scripts' if os.name == 'nt' else 'bin')
142
143    # Delete activation scripts. Typically they're created read-only and venv
144    # will complain when trying to write over them fails.
145    if os.path.isdir(venv_bin):
146        for entry in os.listdir(venv_bin):
147            if entry.lower().startswith('activate'):
148                os.unlink(os.path.join(venv_bin, entry))
149
150    pyvenv_cfg = os.path.join(venv_path, 'pyvenv.cfg')
151    if full_envsetup or not os.path.exists(pyvenv_cfg):
152        # On Mac sometimes the CIPD Python has __PYVENV_LAUNCHER__ set to
153        # point to the system Python, which causes CIPD Python to create
154        # virtualenvs that reference the system Python instead of the CIPD
155        # Python. Clearing __PYVENV_LAUNCHER__ fixes that. See also pwbug/59.
156        envcopy = os.environ.copy()
157        if '__PYVENV_LAUNCHER__' in envcopy:
158            del envcopy['__PYVENV_LAUNCHER__']
159
160        cmd = (python, '-m', 'venv', '--upgrade', venv_path)
161        _check_call(cmd, env=envcopy)
162
163    venv_python = os.path.join(venv_bin, 'python')
164
165    pw_root = os.environ.get('PW_ROOT')
166    if not pw_root and env:
167        pw_root = env.PW_ROOT
168    if not pw_root:
169        pw_root = git_repo_root()
170    if not pw_root:
171        raise GitRepoNotFound()
172
173    # Sometimes we get an error saying "Egg-link ... does not match
174    # installed location". This gets around that. The egg-link files
175    # all come from 'pw'-prefixed packages we installed with --editable.
176    # Source: https://stackoverflow.com/a/48972085
177    for egg_link in glob.glob(
178            os.path.join(venv_path, 'lib/python*/site-packages/*.egg-link')):
179        os.unlink(egg_link)
180
181    def pip_install(*args):
182        cmd = [venv_python, '-m', 'pip', 'install'] + list(args)
183        return _check_call(cmd)
184
185    pip_install('--upgrade', 'pip')
186
187    if requirements:
188        requirement_args = tuple('--requirement={}'.format(req)
189                                 for req in requirements)
190        pip_install('--log', os.path.join(venv_path, 'pip-requirements.log'),
191                    *requirement_args)
192
193    def install_packages(gn_target):
194        if gn_out_dir is None:
195            build_dir = os.path.join(venv_path, gn_target.name)
196        else:
197            build_dir = gn_out_dir
198
199        env_log = 'env-{}.log'.format(gn_target.name)
200        env_log_path = os.path.join(venv_path, env_log)
201        with open(env_log_path, 'w') as outs:
202            for key, value in sorted(os.environ.items()):
203                if key.upper().endswith('PATH'):
204                    print(key, '=', file=outs)
205                    # pylint: disable=invalid-name
206                    for v in value.split(os.pathsep):
207                        print('   ', v, file=outs)
208                    # pylint: enable=invalid-name
209                else:
210                    print(key, '=', value, file=outs)
211
212        gn_log = 'gn-gen-{}.log'.format(gn_target.name)
213        gn_log_path = os.path.join(venv_path, gn_log)
214        try:
215            with open(gn_log_path, 'w') as outs:
216                gn_cmd = (
217                    'gn',
218                    'gen',
219                    build_dir,
220                    '--args=dir_pigweed="{}"'.format(pw_root),
221                )
222                print(gn_cmd, file=outs)
223                subprocess.check_call(gn_cmd,
224                                      cwd=os.path.join(project_root,
225                                                       gn_target.directory),
226                                      stdout=outs,
227                                      stderr=outs)
228        except subprocess.CalledProcessError as err:
229            with open(gn_log_path, 'r') as ins:
230                raise subprocess.CalledProcessError(err.returncode, err.cmd,
231                                                    ins.read())
232
233        ninja_log = 'ninja-{}.log'.format(gn_target.name)
234        ninja_log_path = os.path.join(venv_path, ninja_log)
235        try:
236            with open(ninja_log_path, 'w') as outs:
237                ninja_cmd = ['ninja', '-C', build_dir]
238                ninja_cmd.append(gn_target.target)
239                print(ninja_cmd, file=outs)
240                subprocess.check_call(ninja_cmd, stdout=outs, stderr=outs)
241        except subprocess.CalledProcessError as err:
242            with open(ninja_log_path, 'r') as ins:
243                raise subprocess.CalledProcessError(err.returncode, err.cmd,
244                                                    ins.read())
245
246        with open(os.path.join(venv_path, 'pip-list.log'), 'w') as outs:
247            subprocess.check_call(
248                [venv_python, '-m', 'pip', 'list'],
249                stdout=outs,
250            )
251
252    if gn_targets:
253        if env:
254            env.set('VIRTUAL_ENV', venv_path)
255            env.prepend('PATH', venv_bin)
256            env.clear('PYTHONHOME')
257            with env():
258                for gn_target in gn_targets:
259                    install_packages(gn_target)
260        else:
261            for gn_target in gn_targets:
262                install_packages(gn_target)
263
264    return True
265