[go: up one dir, main page]
Google Git
Sign in
chromium / infra / infra / HEAD / . / go / env.py
blob: b5b187b797f03e20ecf55733435aa409009aa06e [file] [log] [blame]
Takuto Ikutad4400af2021-05-10 08:38:22[diff] [blame]1#!/usr/bin/env vpython3
vadimsh5cfd2fe2014-09-20 00:24:18[diff] [blame]2# Copyright 2014 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
vadimsh5cfd2fe2014-09-20 00:24:18[diff] [blame]5"""Can be used to point environment variable to hermetic Go toolset.
6
7Usage (on linux and mac):
Greg NISBETbb582022023-04-13 18:02:24[diff] [blame]8$ eval "`./env.py`"
vadimsh5cfd2fe2014-09-20 00:24:18[diff] [blame]9$ go version
10
11Or it can be used to wrap a command:
12
13$ ./env.py go version
14"""
15
Vadim Shtayura0cbd3d52021-05-12 22:00:25[diff] [blame]16# This script *must* be the entry point.
vadimsh5cfd2fe2014-09-20 00:24:18[diff] [blame]17assert __name__ == '__main__'
18
vadimsh5cfd2fe2014-09-20 00:24:18[diff] [blame]19import os
20import subprocess
Ricardo Ribaldaea6ea2f2023-11-29 21:41:17[diff] [blame]21import shlex
vadimsh5cfd2fe2014-09-20 00:24:18[diff] [blame]22import sys
23
Vadim Shtayura9b8cf1c2016-10-13 22:29:24[diff] [blame]24
Vadim Shtayura0cbd3d52021-05-12 22:00:25[diff] [blame]25# This should import ./bootstrap.py.
26import bootstrap
27
28# Make sure we picked up the correct bootstrap file.
Andrii Shyshkalov972a2382021-06-06 20:02:02[diff] [blame]29want = os.path.realpath(os.path.abspath(
30 os.path.join(os.path.dirname(__file__), 'bootstrap.py')))
Vadim Shtayuracf100182021-05-13 01:17:35[diff] [blame]31# Allow `bootstrap.__file__` to be a .py or .pyc file.
32if os.path.abspath(bootstrap.__file__) not in (want, want + 'c'):
Vadim Shtayura0cbd3d52021-05-12 22:00:25[diff] [blame]33 raise AssertionError(
34 'Imported wrong bootstrap.py %s instead of %s' %
35 (bootstrap.__file__, want))
vadimsh5cfd2fe2014-09-20 00:24:18[diff] [blame]36
Vadim Shtayura9b8cf1c2016-10-13 22:29:24[diff] [blame]37
dnj665cd3a2016-07-07 17:42:25[diff] [blame]38def _escape_special(v):
39 """Returns (str): The supplied value, with special shell characters escaped.
40
41 Replace special characters with their escaped form. This will allow them
42 to be interpreted by the shell using the $'...' notation.
43
44 Args:
45 v (str): The input value to escape.
46 """
47 for f, r in (
48 ('\n', '\\n'),
49 ('\b', '\\b'),
50 ('\r', '\\r'),
51 ('\t', '\\t'),
52 ('\v', '\\v')):
53 v = v.replace(f, r)
54 return v
55
dnj665cd3a2016-07-07 17:42:25[diff] [blame]56
Vadim Shtayura9b8cf1c2016-10-13 22:29:24[diff] [blame]57if sys.platform == 'win32':
58 def emit_env_var(key, value):
59 # TODO: The quoting here is probably insufficient for all corner cases.
60 # We strip "'" because cmd.exe doesn't like it in PATH for some reason.
Ricardo Ribaldaea6ea2f2023-11-29 21:41:17[diff] [blame]61 print('set %s=%s' % (key, shlex.quote(value).strip("'")))
Vadim Shtayura5cf16bd2021-07-09 17:16:58[diff] [blame]62 def unset_env_var(key):
63 print('set %s=' % (key,))
vadimsh5cfd2fe2014-09-20 00:24:18[diff] [blame]64else:
Vadim Shtayura9b8cf1c2016-10-13 22:29:24[diff] [blame]65 def emit_env_var(key, value):
66 orig_value, value = value, _escape_special(value)
67 # We will only use the $'...' notation if there was an escaped character
68 # in the string.
Vadim Shtayura221a4802021-07-10 00:38:23[diff] [blame]69 print('export %s=%s%s;' % (key, ('$') if orig_value != value else
Ricardo Ribaldaea6ea2f2023-11-29 21:41:17[diff] [blame]70 (''), shlex.quote(value)))
Vadim Shtayura5cf16bd2021-07-09 17:16:58[diff] [blame]71 def unset_env_var(key):
Vadim Shtayura221a4802021-07-10 00:38:23[diff] [blame]72 print('unset %s;' % (key,))
Vadim Shtayura9b8cf1c2016-10-13 22:29:24[diff] [blame]73
74
75def main():
Vadim Shtayura567469b2019-05-23 22:28:55[diff] [blame]76 args = sys.argv[1:]
77 if args and args[0] == '--':
78 args.pop(0)
Dan Jacquesd2f52122017-03-30 00:03:08[diff] [blame]79
Vadim Shtayura5cf16bd2021-07-09 17:16:58[diff] [blame]80 old = os.environ.copy()
Vadim Shtayura567469b2019-05-23 22:28:55[diff] [blame]81 new = bootstrap.prepare_go_environ()
Vadim Shtayura5cf16bd2021-07-09 17:16:58[diff] [blame]82
83 # If given args, it means env.py is executed as a wrapper.
84 if args:
Vadim Shtayura567469b2019-05-23 22:28:55[diff] [blame]85 exe = args[0]
Vadim Shtayura9b8cf1c2016-10-13 22:29:24[diff] [blame]86 if exe == 'python':
87 exe = sys.executable
88 else:
89 # Help Windows to find the executable in new PATH, do it only when
90 # executable is referenced by name (and not by path).
91 if os.sep not in exe:
92 exe = bootstrap.find_executable(exe, [bootstrap.WORKSPACE])
Vadim Shtayura567469b2019-05-23 22:28:55[diff] [blame]93 sys.exit(subprocess.call([exe] + args[1:], env=new))
Vadim Shtayura9b8cf1c2016-10-13 22:29:24[diff] [blame]94
Vadim Shtayura5cf16bd2021-07-09 17:16:58[diff] [blame]95 # If not given args, emit a shell script to modify environ.
96 for key, value in sorted(new.items()):
97 if old.get(key) != value:
98 emit_env_var(key, value)
99 for key in sorted(old):
100 if key not in new:
101 unset_env_var(key)
102
103 # VIRTUAL_ENV is added by the vpython wrapper. It usually *does not* exist
Greg NISBETbb582022023-04-13 18:02:24[diff] [blame]104 # in os.environ of the outer shell that executes eval "`./env.py`". Since we
Vadim Shtayura5cf16bd2021-07-09 17:16:58[diff] [blame]105 # are about to replace the native python in PATH with virtualenv's one, we
106 # must also make sure the new environment has VIRTUAL_ENV set. Otherwise
107 # some virtualenv-aware tools (like gcloud) get confused.
108 #
109 # Note that once env.py finishes execution, nothing is holding a lock on
110 # vpython virtualenv directory, and it may eventually be garbage collected
111 # (while the user is still inside a shell that references it). We assume it
112 # is rare, and the users can manually recover (by reexecuting env.py). This
Greg NISBETbb582022023-04-13 18:02:24[diff] [blame]113 # won't be happening on bots, since they don't use eval "`./env.py`".
Vadim Shtayura5cf16bd2021-07-09 17:16:58[diff] [blame]114 if 'VIRTUAL_ENV' in old:
115 emit_env_var('VIRTUAL_ENV', old['VIRTUAL_ENV'])
116
117 # Warn for common misuse.
118 if sys.platform != 'win32' and sys.stdout.isatty():
119 print()
120 print('# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
121 print('# WRAP THIS COMMAND IN "eval" TO HAVE AN EFFECT!')
Greg NISBETbb582022023-04-13 18:02:24[diff] [blame]122 print('# eval "`./env.py`"')
Vadim Shtayura5cf16bd2021-07-09 17:16:58[diff] [blame]123 print('# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
124
Vadim Shtayura9b8cf1c2016-10-13 22:29:24[diff] [blame]125
Dan Jacquesd2f52122017-03-30 00:03:08[diff] [blame]126main()