1 # Copyright 2020 The Pigweed Authors
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
7 # https://www.apache.org/licenses/LICENSE-2.0
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
14 """Stores the environment changes necessary for Pigweed."""
21 # The order here is important. On Python 2 we want StringIO.StringIO and not
22 # io.StringIO. On Python 3 there is no StringIO module so we want io.StringIO.
23 # Not using six because six is not a standard package we can expect to have
24 # installed in the system Python.
26 from StringIO import StringIO # type: ignore
28 from io import StringIO
30 # Disable super() warnings since this file must be Python 2 compatible.
31 # pylint: disable=super-with-arguments
33 # goto label written to the end of Windows batch files for exiting a script.
34 _SCRIPT_END_LABEL = '_pw_end'
37 class BadNameType(TypeError):
41 class BadValueType(TypeError):
45 class EmptyValue(ValueError):
49 class NewlineInValue(TypeError):
53 class BadVariableName(ValueError):
57 class UnexpectedAction(ValueError):
61 class _Action(object): # pylint: disable=useless-object-inheritance
62 def unapply(self, env, orig_env):
68 def write_deactivate(self,
70 windows=(os.name == 'nt'),
75 class _VariableAction(_Action):
76 # pylint: disable=keyword-arg-before-vararg
77 def __init__(self, name, value, allow_empty_values=False, *args, **kwargs):
78 super(_VariableAction, self).__init__(*args, **kwargs)
81 self.allow_empty_values = allow_empty_values
87 # In python2, unicode is a distinct type.
88 valid_types = (str, unicode)
92 if not isinstance(self.name, valid_types):
93 raise BadNameType('variable name {!r} not of type str'.format(
95 if not isinstance(self.value, valid_types):
96 raise BadValueType('{!r} value {!r} not of type str'.format(
97 self.name, self.value))
99 # Empty strings as environment variable values have different behavior
100 # on different operating systems. Just don't allow them.
101 if not self.allow_empty_values and self.value == '':
102 raise EmptyValue('{!r} value {!r} is the empty string'.format(
103 self.name, self.value))
105 # Many tools have issues with newlines in environment variable values.
106 # Just don't allow them.
107 if '\n' in self.value:
108 raise NewlineInValue('{!r} value {!r} contains a newline'.format(
109 self.name, self.value))
111 if not re.match(r'^[A-Z_][A-Z0-9_]*$', self.name, re.IGNORECASE):
112 raise BadVariableName('bad variable name {!r}'.format(self.name))
114 def unapply(self, env, orig_env):
115 if self.name in orig_env:
116 env[self.name] = orig_env[self.name]
118 env.pop(self.name, None)
121 def _var_form(variable, windows=(os.name == 'nt')):
123 return '%{}%'.format(variable)
124 return '${}'.format(variable)
127 class Set(_VariableAction):
128 """Set a variable."""
129 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
131 for var, replacement in replacements:
133 value = value.replace(replacement, _var_form(var, windows))
136 outs.write('set {name}={value}\n'.format(name=self.name,
139 outs.write('{name}="{value}"\nexport {name}\n'.format(
140 name=self.name, value=value))
142 def write_deactivate(self,
144 windows=(os.name == 'nt'),
146 del replacements # Unused.
149 outs.write('set {name}=\n'.format(name=self.name))
151 outs.write('unset {name}\n'.format(name=self.name))
153 def apply(self, env):
154 env[self.name] = self.value
156 def json(self, data):
157 data['set'][self.name] = self.value
160 class Clear(_VariableAction):
161 """Remove a variable from the environment."""
162 def __init__(self, *args, **kwargs):
164 kwargs['allow_empty_values'] = True
165 super(Clear, self).__init__(*args, **kwargs)
167 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
168 del replacements # Unused.
170 outs.write('set {name}=\n'.format(**vars(self)))
172 outs.write('unset {name}\n'.format(**vars(self)))
174 def apply(self, env):
178 def json(self, data):
179 data['set'][self.name] = None
182 def _initialize_path_like_variable(data, name):
183 default = {'append': [], 'prepend': [], 'remove': []}
184 data['modify'].setdefault(name, default)
187 def _remove_value_from_path(variable, value, pathsep):
188 return ('{variable}="$(echo "${variable}"'
189 ' | sed "s|{pathsep}{value}{pathsep}|{pathsep}|g;"'
190 ' | sed "s|^{value}{pathsep}||g;"'
191 ' | sed "s|{pathsep}{value}$||g;"'
192 ')"\nexport {variable}\n'.format(variable=variable,
197 class Remove(_VariableAction):
198 """Remove a value from a PATH-like variable."""
199 def __init__(self, name, value, pathsep, *args, **kwargs):
200 super(Remove, self).__init__(name, value, *args, **kwargs)
201 self._pathsep = pathsep
203 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
205 for var, replacement in replacements:
207 value = value.replace(replacement, _var_form(var, windows))
211 # TODO(pwbug/231) This does not seem to be supported when value
212 # contains a %variable%. Disabling for now.
213 # outs.write(':: Remove\n:: {value}\n:: from\n:: {name}\n'
214 # ':: before adding it back.\n'
215 # 'set {name}=%{name}:{value}{pathsep}=%\n'.format(
216 # name=self.name, value=value, pathsep=self._pathsep))
219 outs.write('# Remove \n# {value}\n# from\n# {value}\n# before '
221 outs.write(_remove_value_from_path(self.name, value,
224 def apply(self, env):
225 env[self.name] = env[self.name].replace(
226 '{}{}'.format(self.value, self._pathsep), '')
227 env[self.name] = env[self.name].replace(
228 '{}{}'.format(self._pathsep, self.value), '')
230 def json(self, data):
231 _initialize_path_like_variable(data, self.name)
232 data['modify'][self.name]['remove'].append(self.value)
233 if self.value in data['modify'][self.name]['append']:
234 data['modify'][self.name]['append'].remove(self.value)
235 if self.value in data['modify'][self.name]['prepend']:
236 data['modify'][self.name]['prepend'].remove(self.value)
239 class BadVariableValue(ValueError):
243 def _append_prepend_check(action):
244 if '=' in action.value:
245 raise BadVariableValue('"{}" contains "="'.format(action.value))
248 class Prepend(_VariableAction):
249 """Prepend a value to a PATH-like variable."""
250 def __init__(self, name, value, join, *args, **kwargs):
251 super(Prepend, self).__init__(name, value, *args, **kwargs)
254 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
256 for var, replacement in replacements:
258 value = value.replace(replacement, _var_form(var, windows))
259 value = self._join(value, _var_form(self.name, windows))
262 outs.write('set {name}={value}\n'.format(name=self.name,
265 outs.write('{name}="{value}"\nexport {name}\n'.format(
266 name=self.name, value=value))
268 def write_deactivate(self,
270 windows=(os.name == 'nt'),
273 for var, replacement in replacements:
275 value = value.replace(replacement, _var_form(var, windows))
278 _remove_value_from_path(self.name, value, self._join.pathsep))
280 def apply(self, env):
281 env[self.name] = self._join(self.value, env.get(self.name, ''))
284 super(Prepend, self)._check()
285 _append_prepend_check(self)
287 def json(self, data):
288 _initialize_path_like_variable(data, self.name)
289 data['modify'][self.name]['prepend'].append(self.value)
290 if self.value in data['modify'][self.name]['remove']:
291 data['modify'][self.name]['remove'].remove(self.value)
294 class Append(_VariableAction):
295 """Append a value to a PATH-like variable. (Uncommon, see Prepend.)"""
296 def __init__(self, name, value, join, *args, **kwargs):
297 super(Append, self).__init__(name, value, *args, **kwargs)
300 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
302 for var, repl_value in replacements:
304 value = value.replace(repl_value, _var_form(var, windows))
305 value = self._join(_var_form(self.name, windows), value)
308 outs.write('set {name}={value}\n'.format(name=self.name,
311 outs.write('{name}="{value}"\nexport {name}\n'.format(
312 name=self.name, value=value))
314 def write_deactivate(self,
316 windows=(os.name == 'nt'),
319 for var, replacement in replacements:
321 value = value.replace(replacement, _var_form(var, windows))
324 _remove_value_from_path(self.name, value, self._join.pathsep))
326 def apply(self, env):
327 env[self.name] = self._join(env.get(self.name, ''), self.value)
330 super(Append, self)._check()
331 _append_prepend_check(self)
333 def json(self, data):
334 _initialize_path_like_variable(data, self.name)
335 data['modify'][self.name]['append'].append(self.value)
336 if self.value in data['modify'][self.name]['remove']:
337 data['modify'][self.name]['remove'].remove(self.value)
340 class BadEchoValue(ValueError):
345 """Echo a value to the terminal."""
346 def __init__(self, value, newline, *args, **kwargs):
347 # These values act funny on Windows.
348 if value.lower() in ('off', 'on'):
349 raise BadEchoValue(value)
350 super(Echo, self).__init__(*args, **kwargs)
352 self._newline = newline
354 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
355 del replacements # Unused.
356 # POSIX shells parse arguments and pass to echo, but Windows seems to
357 # pass the command line as is without parsing, so quoting is wrong.
361 outs.write('echo.\n')
363 outs.write('echo {}\n'.format(self.value))
365 outs.write('<nul set /p="{}"\n'.format(self.value))
367 # TODO(mohrr) use shlex.quote().
368 outs.write('if [ -z "${PW_ENVSETUP_QUIET:-}" ]; then\n')
370 outs.write(' echo "{}"\n'.format(self.value))
372 outs.write(' echo -n "{}"\n'.format(self.value))
375 def apply(self, env):
379 class Comment(_Action):
380 """Add a comment to the init script."""
381 def __init__(self, value, *args, **kwargs):
382 super(Comment, self).__init__(*args, **kwargs)
385 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
386 del replacements # Unused.
387 comment_char = '::' if windows else '#'
388 for line in self.value.splitlines():
389 outs.write('{} {}\n'.format(comment_char, line))
391 def apply(self, env):
395 class Command(_Action):
397 def __init__(self, command, *args, **kwargs):
398 exit_on_error = kwargs.pop('exit_on_error', True)
399 super(Command, self).__init__(*args, **kwargs)
400 assert isinstance(command, (list, tuple))
401 self.command = command
402 self.exit_on_error = exit_on_error
404 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
405 del replacements # Unused.
406 # TODO(mohrr) use shlex.quote here?
407 outs.write('{}\n'.format(' '.join(self.command)))
408 if not self.exit_on_error:
413 'if %ERRORLEVEL% neq 0 goto {}\n'.format(_SCRIPT_END_LABEL))
415 # Assume failing command produced relevant output.
416 outs.write('if [ "$?" -ne 0 ]; then\n return 1\nfi\n')
418 def apply(self, env):
422 class Doctor(Command):
423 def __init__(self, *args, **kwargs):
424 log_level = 'warn' if 'PW_ENVSETUP_QUIET' in os.environ else 'info'
425 super(Doctor, self).__init__(
426 command=['pw', '--no-banner', '--loglevel', log_level, 'doctor'],
430 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
431 super_call = lambda: super(Doctor, self).write(
432 outs, windows=windows, replacements=replacements)
435 outs.write('if "%PW_ACTIVATE_SKIP_CHECKS%"=="" (\n')
437 outs.write(') else (\n')
438 outs.write('echo Skipping environment check because '
439 'PW_ACTIVATE_SKIP_CHECKS is set\n')
442 outs.write('if [ -z "$PW_ACTIVATE_SKIP_CHECKS" ]; then\n')
445 outs.write('echo Skipping environment check because '
446 'PW_ACTIVATE_SKIP_CHECKS is set\n')
450 class BlankLine(_Action):
451 """Write a blank line to the init script."""
452 def write( # pylint: disable=no-self-use
455 windows=(os.name == 'nt'),
457 del replacements, windows # Unused.
460 def apply(self, env):
464 class Function(_Action):
465 def __init__(self, name, body, *args, **kwargs):
466 super(Function, self).__init__(*args, **kwargs)
470 def write(self, outs, windows=(os.name == 'nt'), replacements=()):
471 del replacements # Unused.
479 """.strip().format(name=self._name, body=self._body))
481 def apply(self, env):
486 def write( # pylint: disable=no-self-use
489 windows=(os.name == 'nt'),
491 del replacements # Unused.
497 # This should detect bash and zsh, which have a hash command that must be
498 # called to get it to forget past commands. Without forgetting past
499 # commands the $PATH changes we made may not be respected.
500 if [ -n "${BASH:-}" -o -n "${ZSH_VERSION:-}" ] ; then
505 def apply(self, env):
509 class Join(object): # pylint: disable=useless-object-inheritance
510 def __init__(self, pathsep=os.pathsep):
511 self.pathsep = pathsep
513 def __call__(self, *args):
514 if len(args) == 1 and isinstance(args[0], (list, tuple)):
516 return self.pathsep.join(args)
519 # TODO(mohrr) remove disable=useless-object-inheritance once in Python 3.
520 # pylint: disable=useless-object-inheritance
521 class Environment(object):
522 """Stores the environment changes necessary for Pigweed.
524 These changes can be accessed by writing them to a file for bash-like
525 shells to source or by using this as a context manager.
527 def __init__(self, *args, **kwargs):
528 pathsep = kwargs.pop('pathsep', os.pathsep)
529 windows = kwargs.pop('windows', os.name == 'nt')
530 allcaps = kwargs.pop('allcaps', windows)
531 super(Environment, self).__init__(*args, **kwargs)
533 self._pathsep = pathsep
534 self._windows = windows
535 self._allcaps = allcaps
536 self._replacements = []
537 self._join = Join(pathsep)
538 self._finalized = False
540 def add_replacement(self, variable, value=None):
541 self._replacements.append((variable, value))
543 def normalize_key(self, name):
547 except AttributeError:
548 # The _Action class has code to handle incorrect types, so
549 # we just ignore this error here.
553 # A newline is printed after each high-level operation. Top-level
554 # operations should not invoke each other (this is why _remove() exists).
556 def set(self, name, value):
557 """Set a variable."""
558 assert not self._finalized
559 name = self.normalize_key(name)
560 self._actions.append(Set(name, value))
563 def clear(self, name):
564 """Remove a variable."""
565 assert not self._finalized
566 name = self.normalize_key(name)
567 self._actions.append(Clear(name))
570 def _remove(self, name, value):
571 """Remove a value from a variable."""
572 assert not self._finalized
573 name = self.normalize_key(name)
574 if self.get(name, None):
575 self._actions.append(Remove(name, value, self._pathsep))
577 def remove(self, name, value):
578 """Remove a value from a PATH-like variable."""
579 assert not self._finalized
580 self._remove(name, value)
583 def append(self, name, value):
584 """Add a value to a PATH-like variable. Rarely used, see prepend()."""
585 assert not self._finalized
586 name = self.normalize_key(name)
587 if self.get(name, None):
588 self._remove(name, value)
589 self._actions.append(Append(name, value, self._join))
591 self._actions.append(Set(name, value))
594 def prepend(self, name, value):
595 """Add a value to the beginning of a PATH-like variable."""
596 assert not self._finalized
597 name = self.normalize_key(name)
598 if self.get(name, None):
599 self._remove(name, value)
600 self._actions.append(Prepend(name, value, self._join))
602 self._actions.append(Set(name, value))
605 def echo(self, value='', newline=True):
606 """Echo a value to the terminal."""
607 # echo() deliberately ignores self._finalized.
608 self._actions.append(Echo(value, newline))
612 def comment(self, comment):
613 """Add a comment to the init script."""
614 # comment() deliberately ignores self._finalized.
615 self._actions.append(Comment(comment))
618 def command(self, command, exit_on_error=True):
620 # command() deliberately ignores self._finalized.
621 self._actions.append(Command(command, exit_on_error=exit_on_error))
625 """Run 'pw doctor'."""
626 self._actions.append(Doctor())
628 def function(self, name, body):
629 """Define a function."""
630 assert not self._finalized
631 self._actions.append(Command(name, body))
634 def _blankline(self):
635 self._actions.append(BlankLine())
638 """Run cleanup at the end of environment setup."""
639 assert not self._finalized
640 self._finalized = True
641 self._actions.append(Hash())
644 if not self._windows:
646 for action in self._actions:
647 action.write_deactivate(buf, windows=self._windows)
648 self._actions.append(Function('_pw_deactivate', buf.getvalue()))
651 def write(self, outs):
652 """Writes a shell init script to outs."""
654 outs.write('@echo off\n')
656 # This is a tuple and not a dictionary because we don't need random
657 # access and order needs to be preserved.
658 replacements = tuple((key, self.get(key) if value is None else value)
659 for key, value in self._replacements)
661 for action in self._actions:
663 windows=self._windows,
664 replacements=replacements)
667 outs.write(':{}\n'.format(_SCRIPT_END_LABEL))
669 def json(self, outs):
675 for action in self._actions:
678 json.dump(data, outs, indent=4, separators=(',', ': '))
681 def write_deactivate(self, outs):
683 outs.write('@echo off\n')
685 for action in reversed(self._actions):
686 action.write_deactivate(outs,
687 windows=self._windows,
690 @contextlib.contextmanager
691 def __call__(self, export=True):
692 """Set environment as if this was written to a file and sourced.
694 Within this context os.environ is updated with the environment
695 defined by this object. If export is False, os.environ is not updated,
696 but in both cases the updated environment is yielded.
698 On exit, previous environment is restored. See contextlib documentation
699 for details on how this function is structured.
702 export(bool): modify the environment of the running process (and
703 thus, its subprocesses)
705 Yields the new environment object.
709 orig_env = os.environ.copy()
712 env = os.environ.copy()
714 for action in self._actions:
721 for action in self._actions:
722 action.unapply(env=os.environ, orig_env=orig_env)
724 def get(self, key, default=None):
725 """Get the value of a variable within context of this object."""
726 key = self.normalize_key(key)
727 with self(export=False) as env:
728 return env.get(key, default)
730 def __getitem__(self, key):
731 """Get the value of a variable within context of this object."""
732 key = self.normalize_key(key)
733 with self(export=False) as env: