• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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"""Stores the environment changes necessary for Pigweed."""
15
16import contextlib
17import os
18import re
19
20# The order here is important. On Python 2 we want StringIO.StringIO and not
21# io.StringIO. On Python 3 there is no StringIO module so we want io.StringIO.
22# Not using six because six is not a standard package we can expect to have
23# installed in the system Python.
24try:
25    from StringIO import StringIO  # type: ignore
26except ImportError:
27    from io import StringIO
28
29from . import apply_visitor
30from . import batch_visitor
31from . import gni_visitor
32from . import json_visitor
33from . import shell_visitor
34
35# Disable super() warnings since this file must be Python 2 compatible.
36# pylint: disable=super-with-arguments
37
38
39class BadNameType(TypeError):
40    pass
41
42
43class BadValueType(TypeError):
44    pass
45
46
47class EmptyValue(ValueError):
48    pass
49
50
51class NewlineInValue(TypeError):
52    pass
53
54
55class BadVariableName(ValueError):
56    pass
57
58
59class UnexpectedAction(ValueError):
60    pass
61
62
63class AcceptNotOverridden(TypeError):
64    pass
65
66
67class _Action(object):  # pylint: disable=useless-object-inheritance
68    def unapply(self, env, orig_env):
69        pass
70
71    def accept(self, visitor):
72        del visitor
73        raise AcceptNotOverridden(
74            'accept() not overridden for {}'.format(self.__class__.__name__)
75        )
76
77    def write_deactivate(
78        self, outs, windows=(os.name == 'nt'), replacements=()
79    ):
80        pass
81
82
83class _VariableAction(_Action):
84    # pylint: disable=keyword-arg-before-vararg
85    def __init__(self, name, value, allow_empty_values=False, *args, **kwargs):
86        super(_VariableAction, self).__init__(*args, **kwargs)
87        self.name = name
88        self.value = value
89        self.allow_empty_values = allow_empty_values
90
91        self._check()
92
93    def _check(self):
94        try:
95            # In python2, unicode is a distinct type.
96            valid_types = (str, unicode)
97        except NameError:
98            valid_types = (str,)
99
100        if not isinstance(self.name, valid_types):
101            raise BadNameType(
102                'variable name {!r} not of type str'.format(self.name)
103            )
104        if not isinstance(self.value, valid_types):
105            raise BadValueType(
106                '{!r} value {!r} not of type str'.format(self.name, self.value)
107            )
108
109        # Empty strings as environment variable values have different behavior
110        # on different operating systems. Just don't allow them.
111        if not self.allow_empty_values and self.value == '':
112            raise EmptyValue(
113                '{!r} value {!r} is the empty string'.format(
114                    self.name, self.value
115                )
116            )
117
118        # Many tools have issues with newlines in environment variable values.
119        # Just don't allow them.
120        if '\n' in self.value:
121            raise NewlineInValue(
122                '{!r} value {!r} contains a newline'.format(
123                    self.name, self.value
124                )
125            )
126
127        if not re.match(r'^[A-Z_][A-Z0-9_]*$', self.name, re.IGNORECASE):
128            raise BadVariableName('bad variable name {!r}'.format(self.name))
129
130    def unapply(self, env, orig_env):
131        if self.name in orig_env:
132            env[self.name] = orig_env[self.name]
133        else:
134            env.pop(self.name, None)
135
136    def __repr__(self):
137        return '{}({}, {})'.format(
138            self.__class__.__name__, self.name, self.value
139        )
140
141
142class Set(_VariableAction):
143    """Set a variable."""
144
145    def __init__(self, *args, **kwargs):
146        deactivate = kwargs.pop('deactivate', True)
147        super(Set, self).__init__(*args, **kwargs)
148        self.deactivate = deactivate
149
150    def accept(self, visitor):
151        visitor.visit_set(self)
152
153
154class Clear(_VariableAction):
155    """Remove a variable from the environment."""
156
157    def __init__(self, *args, **kwargs):
158        kwargs['value'] = ''
159        kwargs['allow_empty_values'] = True
160        super(Clear, self).__init__(*args, **kwargs)
161
162    def accept(self, visitor):
163        visitor.visit_clear(self)
164
165
166class Remove(_VariableAction):
167    """Remove a value from a PATH-like variable."""
168
169    def accept(self, visitor):
170        visitor.visit_remove(self)
171
172
173class BadVariableValue(ValueError):
174    pass
175
176
177def _append_prepend_check(action):
178    if '=' in action.value:
179        raise BadVariableValue('"{}" contains "="'.format(action.value))
180
181
182class Prepend(_VariableAction):
183    """Prepend a value to a PATH-like variable."""
184
185    def __init__(self, name, value, join, *args, **kwargs):
186        super(Prepend, self).__init__(name, value, *args, **kwargs)
187        self._join = join
188
189    def _check(self):
190        super(Prepend, self)._check()
191        _append_prepend_check(self)
192
193    def accept(self, visitor):
194        visitor.visit_prepend(self)
195
196
197class Append(_VariableAction):
198    """Append a value to a PATH-like variable. (Uncommon, see Prepend.)"""
199
200    def __init__(self, name, value, join, *args, **kwargs):
201        super(Append, self).__init__(name, value, *args, **kwargs)
202        self._join = join
203
204    def _check(self):
205        super(Append, self)._check()
206        _append_prepend_check(self)
207
208    def accept(self, visitor):
209        visitor.visit_append(self)
210
211
212class BadEchoValue(ValueError):
213    pass
214
215
216class Echo(_Action):
217    """Echo a value to the terminal."""
218
219    def __init__(self, value, newline, *args, **kwargs):
220        # These values act funny on Windows.
221        if value.lower() in ('off', 'on'):
222            raise BadEchoValue(value)
223        super(Echo, self).__init__(*args, **kwargs)
224        self.value = value
225        self.newline = newline
226
227    def accept(self, visitor):
228        visitor.visit_echo(self)
229
230    def __repr__(self):
231        return 'Echo({}, newline={})'.format(self.value, self.newline)
232
233
234class Comment(_Action):
235    """Add a comment to the init script."""
236
237    def __init__(self, value, *args, **kwargs):
238        super(Comment, self).__init__(*args, **kwargs)
239        self.value = value
240
241    def accept(self, visitor):
242        visitor.visit_comment(self)
243
244    def __repr__(self):
245        return 'Comment({})'.format(self.value)
246
247
248class Command(_Action):
249    """Run a command."""
250
251    def __init__(self, command, *args, **kwargs):
252        exit_on_error = kwargs.pop('exit_on_error', True)
253        super(Command, self).__init__(*args, **kwargs)
254        assert isinstance(command, (list, tuple))
255        self.command = command
256        self.exit_on_error = exit_on_error
257
258    def accept(self, visitor):
259        visitor.visit_command(self)
260
261    def __repr__(self):
262        return 'Command({})'.format(self.command)
263
264
265class Doctor(Command):
266    def __init__(self, *args, **kwargs):
267        log_level = 'warn' if 'PW_ENVSETUP_QUIET' in os.environ else 'info'
268        cmd = ['pw', '--no-banner', '--loglevel', log_level, 'doctor']
269        super(Doctor, self).__init__(command=cmd, *args, **kwargs)
270
271    def accept(self, visitor):
272        visitor.visit_doctor(self)
273
274    def __repr__(self):
275        return 'Doctor()'
276
277
278class BlankLine(_Action):
279    """Write a blank line to the init script."""
280
281    def accept(self, visitor):
282        visitor.visit_blank_line(self)
283
284    def __repr__(self):
285        return 'BlankLine()'
286
287
288class Function(_Action):
289    def __init__(self, name, body, *args, **kwargs):
290        super(Function, self).__init__(*args, **kwargs)
291        self.name = name
292        self.body = body
293
294    def accept(self, visitor):
295        visitor.visit_function(self)
296
297    def __repr__(self):
298        return 'Function({}, {})'.format(self.name, self.body)
299
300
301class Hash(_Action):
302    def accept(self, visitor):
303        visitor.visit_hash(self)
304
305    def __repr__(self):
306        return 'Hash()'
307
308
309class Join(object):  # pylint: disable=useless-object-inheritance
310    def __init__(self, pathsep=os.pathsep):
311        self.pathsep = pathsep
312
313
314# TODO(mohrr) remove disable=useless-object-inheritance once in Python 3.
315# pylint: disable=useless-object-inheritance
316class Environment(object):
317    """Stores the environment changes necessary for Pigweed.
318
319    These changes can be accessed by writing them to a file for bash-like
320    shells to source or by using this as a context manager.
321    """
322
323    def __init__(self, *args, **kwargs):
324        pathsep = kwargs.pop('pathsep', os.pathsep)
325        windows = kwargs.pop('windows', os.name == 'nt')
326        allcaps = kwargs.pop('allcaps', windows)
327        super(Environment, self).__init__(*args, **kwargs)
328        self._actions = []
329        self._pathsep = pathsep
330        self._windows = windows
331        self._allcaps = allcaps
332        self.replacements = []
333        self._join = Join(pathsep)
334        self._finalized = False
335
336    def add_replacement(self, variable, value=None):
337        self.replacements.append((variable, value))
338
339    def normalize_key(self, name):
340        if self._allcaps:
341            try:
342                return name.upper()
343            except AttributeError:
344                # The _Action class has code to handle incorrect types, so
345                # we just ignore this error here.
346                pass
347        return name
348
349    # A newline is printed after each high-level operation. Top-level
350    # operations should not invoke each other (this is why _remove() exists).
351
352    def set(self, name, value, deactivate=True):
353        """Set a variable."""
354        assert not self._finalized
355        name = self.normalize_key(name)
356        self._actions.append(Set(name, value, deactivate=deactivate))
357        self._blankline()
358
359    def clear(self, name):
360        """Remove a variable."""
361        assert not self._finalized
362        name = self.normalize_key(name)
363        self._actions.append(Clear(name))
364        self._blankline()
365
366    def _remove(self, name, value):
367        """Remove a value from a variable."""
368        assert not self._finalized
369        name = self.normalize_key(name)
370        if self.get(name, None):
371            self._actions.append(Remove(name, value, self._pathsep))
372
373    def remove(self, name, value):
374        """Remove a value from a PATH-like variable."""
375        assert not self._finalized
376        self._remove(name, value)
377        self._blankline()
378
379    def append(self, name, value):
380        """Add a value to a PATH-like variable. Rarely used, see prepend()."""
381        assert not self._finalized
382        name = self.normalize_key(name)
383        if self.get(name, None):
384            self._remove(name, value)
385            self._actions.append(Append(name, value, self._join))
386        else:
387            self._actions.append(Set(name, value))
388        self._blankline()
389
390    def prepend(self, name, value):
391        """Add a value to the beginning of a PATH-like variable."""
392        assert not self._finalized
393        name = self.normalize_key(name)
394        if self.get(name, None):
395            self._remove(name, value)
396            self._actions.append(Prepend(name, value, self._join))
397        else:
398            self._actions.append(Set(name, value))
399        self._blankline()
400
401    def echo(self, value='', newline=True):
402        """Echo a value to the terminal."""
403        # echo() deliberately ignores self._finalized.
404        self._actions.append(Echo(value, newline))
405        if value:
406            self._blankline()
407
408    def comment(self, comment):
409        """Add a comment to the init script."""
410        # comment() deliberately ignores self._finalized.
411        self._actions.append(Comment(comment))
412        self._blankline()
413
414    def command(self, command, exit_on_error=True):
415        """Run a command."""
416        # command() deliberately ignores self._finalized.
417        self._actions.append(Command(command, exit_on_error=exit_on_error))
418        self._blankline()
419
420    def doctor(self):
421        """Run 'pw doctor'."""
422        self._actions.append(Doctor())
423
424    def function(self, name, body):
425        """Define a function."""
426        assert not self._finalized
427        self._actions.append(Command(name, body))
428        self._blankline()
429
430    def _blankline(self):
431        self._actions.append(BlankLine())
432
433    def finalize(self):
434        """Run cleanup at the end of environment setup."""
435        assert not self._finalized
436        self._finalized = True
437        self._actions.append(Hash())
438        self._blankline()
439
440        if not self._windows:
441            buf = StringIO()
442            self.write_deactivate(buf)
443            self._actions.append(Function('_pw_deactivate', buf.getvalue()))
444            self._blankline()
445
446    def accept(self, visitor):
447        for action in self._actions:
448            action.accept(visitor)
449
450    def gni(self, outs, project_root):
451        gni_visitor.GNIVisitor(project_root).serialize(self, outs)
452
453    def json(self, outs):
454        json_visitor.JSONVisitor().serialize(self, outs)
455
456    def write(self, outs):
457        if self._windows:
458            visitor = batch_visitor.BatchVisitor(pathsep=self._pathsep)
459        else:
460            visitor = shell_visitor.ShellVisitor(pathsep=self._pathsep)
461        visitor.serialize(self, outs)
462
463    def write_deactivate(self, outs):
464        if self._windows:
465            return
466        visitor = shell_visitor.DeactivateShellVisitor(pathsep=self._pathsep)
467        visitor.serialize(self, outs)
468
469    @contextlib.contextmanager
470    def __call__(self, export=True):
471        """Set environment as if this was written to a file and sourced.
472
473        Within this context os.environ is updated with the environment
474        defined by this object. If export is False, os.environ is not updated,
475        but in both cases the updated environment is yielded.
476
477        On exit, previous environment is restored. See contextlib documentation
478        for details on how this function is structured.
479
480        Args:
481          export(bool): modify the environment of the running process (and
482            thus, its subprocesses)
483
484        Yields the new environment object.
485        """
486        orig_env = {}
487        try:
488            if export:
489                orig_env = os.environ.copy()
490                env = os.environ
491            else:
492                env = os.environ.copy()
493
494            apply = apply_visitor.ApplyVisitor(pathsep=self._pathsep)
495            apply.apply(self, env)
496
497            yield env
498
499        finally:
500            if export:
501                for key in set(os.environ):
502                    try:
503                        os.environ[key] = orig_env[key]
504                    except KeyError:
505                        del os.environ[key]
506                for key in set(orig_env) - set(os.environ):
507                    os.environ[key] = orig_env[key]
508
509    def get(self, key, default=None):
510        """Get the value of a variable within context of this object."""
511        key = self.normalize_key(key)
512        with self(export=False) as env:
513            return env.get(key, default)
514
515    def __getitem__(self, key):
516        """Get the value of a variable within context of this object."""
517        key = self.normalize_key(key)
518        with self(export=False) as env:
519            return env[key]
520