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