xref: /aosp_15_r20/external/pigweed/pw_env_setup/py/pw_env_setup/shell_visitor.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2021 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"""Serializes an Environment into a shell file."""
15
16import inspect
17
18
19class _BaseShellVisitor:
20    def __init__(self, *args, **kwargs):
21        pathsep = kwargs.pop('pathsep', ':')
22        super().__init__(*args, **kwargs)
23        self._pathsep = pathsep
24        self._outs = None
25
26    def _remove_value_from_path(self, variable, value):
27        return (
28            '{variable}="$(echo "${variable}"'
29            ' | sed "s|{pathsep}{value}{pathsep}|{pathsep}|g;"'
30            ' | sed "s|^{value}{pathsep}||g;"'
31            ' | sed "s|{pathsep}{value}$||g;"'
32            ')"\nexport {variable}\n'.format(
33                variable=variable, value=value, pathsep=self._pathsep
34            )
35        )
36
37    def visit_hash(self, hash):  # pylint: disable=redefined-builtin
38        del hash
39        self._outs.write(
40            inspect.cleandoc(
41                '''
42        # This should detect bash and zsh, which have a hash command that must
43        # be called to get it to forget past commands. Without forgetting past
44        # commands the $PATH changes we made may not be respected.
45        if [ -n "${BASH:-}" -o -n "${ZSH_VERSION:-}" ] ; then
46            hash -r\n
47        fi
48        '''
49            )
50        )
51
52
53class ShellVisitor(_BaseShellVisitor):
54    """Serializes an Environment into a bash-like shell file."""
55
56    def __init__(self, *args, **kwargs):
57        super().__init__(*args, **kwargs)
58        self._replacements = ()
59
60    def serialize(self, env, outs):
61        """Write a shell file based on the given environment.
62
63        Args:
64            env (environment.Environment): Environment variables to use.
65            outs (file): Shell file to write.
66        """
67        try:
68            self._replacements = tuple(
69                (key, env.get(key) if value is None else value)
70                for key, value in env.replacements
71            )
72            self._outs = outs
73
74            env.accept(self)
75
76        finally:
77            self._replacements = ()
78            self._outs = None
79
80    def _apply_replacements(self, action):
81        value = action.value
82        for var, replacement in self._replacements:
83            if var != action.name:
84                value = value.replace(replacement, '${}'.format(var))
85        return value
86
87    def visit_set(self, set):  # pylint: disable=redefined-builtin
88        value = self._apply_replacements(set)
89        self._outs.write(
90            '{name}="{value}"\nexport {name}\n'.format(
91                name=set.name, value=value
92            )
93        )
94
95    def visit_clear(self, clear):
96        self._outs.write('unset {name}\n'.format(**vars(clear)))
97
98    def visit_remove(self, remove):
99        value = self._apply_replacements(remove)
100        self._outs.write(
101            '# Remove \n#   {value}\n# from {name} before adding it '
102            'back.\n'.format(value=remove.value, name=remove.name)
103        )
104        self._outs.write(self._remove_value_from_path(remove.name, value))
105
106    def _join(self, *args):
107        if len(args) == 1 and isinstance(args[0], (list, tuple)):
108            args = args[0]
109        return self._pathsep.join(args)
110
111    def visit_prepend(self, prepend):
112        value = self._apply_replacements(prepend)
113        value = self._join(value, '${}'.format(prepend.name))
114        self._outs.write(
115            '{name}="{value}"\nexport {name}\n'.format(
116                name=prepend.name, value=value
117            )
118        )
119
120    def visit_append(self, append):
121        value = self._apply_replacements(append)
122        value = self._join('${}'.format(append.name), value)
123        self._outs.write(
124            '{name}="{value}"\nexport {name}\n'.format(
125                name=append.name, value=value
126            )
127        )
128
129    def visit_echo(self, echo):
130        # TODO(mohrr) use shlex.quote().
131        self._outs.write('if [ -z "${PW_ENVSETUP_QUIET:-}" ]; then\n')
132        if echo.newline:
133            self._outs.write('  echo "{}"\n'.format(echo.value))
134        else:
135            self._outs.write('  echo -n "{}"\n'.format(echo.value))
136        self._outs.write('fi\n')
137
138    def visit_comment(self, comment):
139        for line in comment.value.splitlines():
140            self._outs.write('# {}\n'.format(line))
141
142    def visit_command(self, command):
143        # TODO(mohrr) use shlex.quote here?
144        self._outs.write('{}\n'.format(' '.join(command.command)))
145        if not command.exit_on_error:
146            return
147
148        # Assume failing command produced relevant output.
149        self._outs.write('if [ "$?" -ne 0 ]; then\n  return 1\nfi\n')
150
151    def visit_doctor(self, doctor):
152        self._outs.write('if [ -z "$PW_ACTIVATE_SKIP_CHECKS" ]; then\n')
153        self.visit_command(doctor)
154        self._outs.write('else\n')
155        self._outs.write(
156            'echo Skipping environment check because '
157            'PW_ACTIVATE_SKIP_CHECKS is set\n'
158        )
159        self._outs.write('fi\n')
160
161    def visit_blank_line(self, blank_line):
162        del blank_line
163        self._outs.write('\n')
164
165    def visit_function(self, function):
166        self._outs.write(
167            '{name}() {{\n{body}\n}}\n'.format(
168                name=function.name, body=function.body
169            )
170        )
171
172
173class DeactivateShellVisitor(_BaseShellVisitor):
174    """Removes values from a bash-like shell environment."""
175
176    def __init__(self, *args, **kwargs):
177        pathsep = kwargs.pop('pathsep', ':')
178        super().__init__(*args, **kwargs)
179        self._pathsep = pathsep
180
181    def serialize(self, env, outs):
182        try:
183            self._outs = outs
184
185            env.accept(self)
186
187        finally:
188            self._outs = None
189
190    def visit_set(self, set):  # pylint: disable=redefined-builtin
191        if set.deactivate:
192            self._outs.write('unset {name}\n'.format(name=set.name))
193
194    def visit_clear(self, clear):
195        pass  # Not relevant.
196
197    def visit_remove(self, remove):
198        pass  # Not relevant.
199
200    def visit_prepend(self, prepend):
201        self._outs.write(
202            self._remove_value_from_path(prepend.name, prepend.value)
203        )
204
205    def visit_append(self, append):
206        self._outs.write(
207            self._remove_value_from_path(append.name, append.value)
208        )
209
210    def visit_echo(self, echo):
211        pass  # Not relevant.
212
213    def visit_comment(self, comment):
214        pass  # Not relevant.
215
216    def visit_command(self, command):
217        pass  # Not relevant.
218
219    def visit_doctor(self, doctor):
220        pass  # Not relevant.
221
222    def visit_blank_line(self, blank_line):
223        pass  # Not relevant.
224
225    def visit_function(self, function):
226        pass  # Not relevant.
227
228
229class FishShellVisitor(ShellVisitor):
230    """Serializes an Environment into a fish shell file."""
231
232    def __init__(self, *args, **kwargs):
233        super().__init__(*args, **kwargs)
234        self._pathsep = ' '
235
236    def _remove_value_from_path(self, variable, value):
237        return 'set PATH (string match -v {value} ${variable})\n'.format(
238            variable=variable, value=value
239        )
240
241    def visit_set(self, set):  # pylint: disable=redefined-builtin
242        value = self._apply_replacements(set)
243        self._outs.write(
244            'set -x {name} {value}\n'.format(name=set.name, value=value)
245        )
246
247    def visit_clear(self, clear):
248        self._outs.write('set -e {name}\n'.format(**vars(clear)))
249
250    def visit_remove(self, remove):
251        value = self._apply_replacements(remove)
252        self._remove_value_from_path(remove.name, value)
253
254    def visit_prepend(self, prepend):
255        value = self._apply_replacements(prepend)
256        self._outs.write(
257            'set -x --prepend {name} {value}\n'.format(
258                name=prepend.name, value=value
259            )
260        )
261
262    def visit_append(self, append):
263        value = self._apply_replacements(append)
264        self._outs.write(
265            'set -x --append {name} {value}\n'.format(
266                name=append.name, value=value
267            )
268        )
269
270    def visit_echo(self, echo):
271        self._outs.write('if not set -q PW_ENVSETUP_QUIET\n')
272        if echo.newline:
273            self._outs.write('  echo "{}"\n'.format(echo.value))
274        else:
275            self._outs.write('  echo -n "{}"\n'.format(echo.value))
276        self._outs.write('end\n')
277
278    def visit_hash(self, hash):  # pylint: disable=redefined-builtin
279        del hash
280
281    def visit_function(self, function):
282        self._outs.write(
283            'function {name}\n{body}\nend\n'.format(
284                name=function.name, body=function.body
285            )
286        )
287
288    def visit_command(self, command):
289        self._outs.write('{}\n'.format(' '.join(command.command)))
290        if not command.exit_on_error:
291            return
292
293        # Assume failing command produced relevant output.
294        self._outs.write('if test $status -ne 0\n  return 1\nend\n')
295
296    def visit_doctor(self, doctor):
297        self._outs.write('if not set -q PW_ACTIVATE_SKIP_CHECKS\n')
298        self.visit_command(doctor)
299        self._outs.write('else\n')
300        self._outs.write(
301            'echo Skipping environment check because '
302            'PW_ACTIVATE_SKIP_CHECKS is set\n'
303        )
304        self._outs.write('end\n')
305
306
307class DeactivateFishShellVisitor(FishShellVisitor):
308    """Removes values from a fish shell environment."""
309
310    def serialize(self, env, outs):
311        try:
312            self._outs = outs
313
314            env.accept(self)
315
316        finally:
317            self._outs = None
318
319    def visit_set(self, set):  # pylint: disable=redefined-builtin
320        if set.deactivate:
321            self._outs.write('set -e {name}\n'.format(name=set.name))
322
323    def visit_clear(self, clear):
324        pass  # Not relevant.
325
326    def visit_remove(self, remove):
327        pass  # Not relevant.
328
329    def visit_prepend(self, prepend):
330        self._outs.write(
331            self._remove_value_from_path(prepend.name, prepend.value)
332        )
333
334    def visit_append(self, append):
335        self._outs.write(
336            self._remove_value_from_path(append.name, append.value)
337        )
338
339    def visit_echo(self, echo):
340        pass  # Not relevant.
341
342    def visit_comment(self, comment):
343        pass  # Not relevant.
344
345    def visit_command(self, command):
346        pass  # Not relevant.
347
348    def visit_doctor(self, doctor):
349        pass  # Not relevant.
350
351    def visit_blank_line(self, blank_line):
352        pass  # Not relevant.
353
354    def visit_function(self, function):
355        pass  # Not relevant.
356