xref: /llvm-project/llvm/utils/TableGen/jupyter/tablegen_kernel/kernel.py (revision f8559751fc2b15b45ac417be9abe865085af45ad)
1# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2# See https://llvm.org/LICENSE.txt for license information.
3# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4
5import os
6import shutil
7import subprocess
8import tempfile
9from ipykernel.kernelbase import Kernel
10
11__version__ = "0.0.1"
12
13
14class TableGenKernelException(Exception):
15    pass
16
17
18class TableGenKernel(Kernel):
19    """Kernel using llvm-tblgen inside jupyter.
20
21    All input is treated as TableGen unless the first non whitespace character
22    is "%" in which case it is a "magic" line.
23
24    The supported cell magic is:
25    * %args    - to set the arguments passed to llvm-tblgen.
26    * %reset   - to reset the cached code and magic state.
27    * %noreset - to not reset the cached code and magic state
28                 (useful when you have changed the default to always
29                  reset the cache).
30
31    These are "cell magic" meaning it applies to the whole cell. Therefore
32    it must be the first line, or part of a run of magic lines starting
33    from the first line.
34
35    The following are global magic (that applies to all cells going
36    forward):
37    * %config  - to change the behaviour of the kernel overall, including
38                 changing defaults for things like resets.
39
40    Global magic must be written in the same way as cell magic.
41
42    ```tablgen
43    %args
44    %reset
45    %args --print-records --print-detailed-records
46    class Stuff {
47      string Name;
48    }
49
50    def a_thing : Stuff {}
51    ```
52
53    """
54
55    implementation = "tablegen"
56    implementation_version = __version__
57
58    language_version = __version__
59    language = "tablegen"
60    language_info = {
61        "name": "tablegen",
62        "mimetype": "text/x-tablegen",
63        "file_extension": ".td",
64        "pygments_lexer": "text",
65    }
66
67    def __init__(self, **kwargs):
68        Kernel.__init__(self, **kwargs)
69        self._executable = None
70        # A list of (code, magic) tuples.
71        # All the previous cell's code we have run since the last reset.
72        # This emulates a persistent state like a Python interpreter would have.
73        self._previous_code = ""
74        # The most recent set of magic since the last reset.
75        self._previous_magic = {}
76        # The default cache reset behaviour. True means do not cache anything
77        # between cells.
78        self._cell_reset = False
79
80    @property
81    def banner(self):
82        return "llvm-tblgen kernel %s" % __version__
83
84    def get_executable(self):
85        """If this is the first run, search for llvm-tblgen.
86        Otherwise return the cached path to it."""
87        if self._executable is None:
88            path = os.environ.get("LLVM_TBLGEN_EXECUTABLE")
89            if path is not None and os.path.isfile(path) and os.access(path, os.X_OK):
90                self._executable = path
91            else:
92                path = shutil.which("llvm-tblgen")
93                if path is None:
94                    raise OSError(
95                        "llvm-tblgen not found. Put it on your PATH or set the"
96                        " environment variable LLVM_TBLGEN_EXECUTABLE to point to it."
97                    )
98                self._executable = path
99
100        return self._executable
101
102    def parse_config_magic(self, config):
103        """Config should be a list of parameters given to the %config command.
104        We allow only one setting per %config line and that setting can only
105        have one value.
106
107        Assuming the parameters are valid, update the kernel's setting with
108        the new value.
109
110        If there is an error, raise a TableGenKernelException.
111
112        >>> k.parse_config_magic([])
113        Traceback (most recent call last):
114         ...
115        TableGenKernelException: Incorrect number of parameters to %config. Expected %config <setting> <value>.
116        >>> k._cell_reset
117        False
118        >>> k.parse_config_magic(["a", "b", "c"])
119        Traceback (most recent call last):
120         ...
121        TableGenKernelException: Incorrect number of parameters to %config. Expected %config <setting> <value>.
122        >>> k.parse_config_magic(["notasetting", "..."])
123        Traceback (most recent call last):
124         ...
125        TableGenKernelException: Unknown kernel setting "notasetting". Possible settings are: "cellreset".
126        >>> k.parse_config_magic(["cellreset", "food"])
127        Traceback (most recent call last):
128         ...
129        TableGenKernelException: Invalid value for setting "cellreset", expected "on" or "off".
130        >>> k.parse_config_magic(["cellreset", "on"])
131        >>> k._cell_reset
132        True
133        >>> k.parse_config_magic(["cellreset", "off"])
134        >>> k._cell_reset
135        False
136        """
137        if len(config) != 2:
138            raise TableGenKernelException(
139                "Incorrect number of parameters to %config. Expected %config <setting> <value>."
140            )
141
142        name, value = config
143        if name != "cellreset":
144            raise TableGenKernelException(
145                'Unknown kernel setting "{}". '
146                'Possible settings are: "cellreset".'.format(name)
147            )
148
149        try:
150            self._cell_reset = {"on": True, "off": False}[value.lower()]
151        except KeyError:
152            raise TableGenKernelException(
153                'Invalid value for setting "{}", '
154                'expected "on" or "off".'.format(name)
155            )
156
157    def get_magic(self, code):
158        """Given a block of code remove the magic lines from it.
159        Returns a tuple of newline joined code lines and a dictionary of magic.
160        Where the key is the magic name (minus the %) and the values are lists
161        of the arguments to the magic.
162
163        Currently we only look for "cell magic" which must be at the start of
164        the cell. Meaning the first line, or a set of lines beginning with %
165        that come before the first non-magic line.
166
167        >>> k.get_magic("")
168        ('', {})
169        >>> k.get_magic("Hello World.\\nHello again.")
170        ('Hello World.\\nHello again.', {})
171        >>> k.get_magic("   %foo a b c")
172        ('', {'foo': ['a', 'b', 'c']})
173        >>> k.get_magic("%foo\\n   %foo a b c\\nFoo")
174        ('Foo', {'foo': ['a', 'b', 'c']})
175        >>> k.get_magic("%foo\\n%bar\\nFoo")
176        ('Foo', {'foo': [], 'bar': []})
177        >>> k.get_magic("Foo\\n%foo\\nFoo")
178        ('Foo\\n%foo\\nFoo', {})
179        >>> k.get_magic("%bar\\n\\n%foo")
180        ('\\n%foo', {'bar': []})
181        >>> k.get_magic("%foo a b\\n   Foo\\n%foo c d")
182        ('   Foo\\n%foo c d', {'foo': ['a', 'b']})
183        >>> k.get_magic("%foo a b\\n \\n%foo c d")
184        (' \\n%foo c d', {'foo': ['a', 'b']})
185        """
186        magic = {}
187        code_lines = []
188
189        lines = code.splitlines()
190        while lines:
191            line = lines.pop(0)
192            possible_magic = line.lstrip()
193            if possible_magic.startswith("%"):
194                magic_parts = possible_magic.split()
195                # Key has the % removed
196                magic[magic_parts[0][1:]] = magic_parts[1:]
197            else:
198                code_lines = [line, *lines]
199                break
200
201        return "\n".join(code_lines), magic
202
203    def should_reset(self, magic):
204        """Return true if we should reset the cache, based on the default
205        setting and the current cell's magic %reset and/or %noreset.
206
207        >>> k._cell_reset = False
208        >>> k.should_reset({})
209        False
210        >>> k.should_reset({'reset': [], 'noreset': []})
211        Traceback (most recent call last):
212        ...
213        TableGenKernelException: %reset and %noreset in the same cell is not allowed. Use only one, or neither.
214        >>> k.should_reset({'reset': []})
215        True
216        >>> k.should_reset({'noreset': []})
217        False
218        >>> k._cell_reset = True
219        >>> k.should_reset({})
220        True
221        >>> k.should_reset({'reset': [], 'noreset': []})
222        Traceback (most recent call last):
223        ...
224        TableGenKernelException: %reset and %noreset in the same cell is not allowed. Use only one, or neither.
225        >>> k.should_reset({'reset': []})
226        True
227        >>> k.should_reset({'noreset': []})
228        False
229        """
230        # Cell reset is the default unless told otherwise.
231        should_reset = self._cell_reset
232        # Magic reset commands always win if present.
233        reset = magic.get("reset") is not None
234        noreset = magic.get("noreset") is not None
235
236        if reset and not noreset:
237            should_reset = True
238        elif noreset and not reset:
239            should_reset = False
240        elif noreset and reset:
241            raise TableGenKernelException(
242                "%reset and %noreset in the same cell is not allowed. Use only one, or neither."
243            )
244        # else neither are set so use the default.
245
246        return should_reset
247
248    def get_code_and_args(self, new_code):
249        """Get the code that do_execute should use, taking into account
250        the code from any cached cells.
251
252        Returns the code to compile and the arguments to use to do so.
253
254        >>> k._previous_code = ""
255        >>> k._previous_magic = {}
256        >>> k.get_code_and_args("")
257        ('', [])
258        >>> k.get_code_and_args("%args 1\\nSome code")
259        ('Some code', ['1'])
260        >>> k.get_code_and_args("%args 2\\nSome more code")
261        ('Some code\\nSome more code', ['2'])
262        >>> k.get_code_and_args("%reset\\n%args 3 4\\nSome new code")
263        ('Some new code', ['3', '4'])
264        >>> k.get_code_and_args("%reset\\nSome new code")
265        ('Some new code', [])
266        """
267        new_code, new_magic = self.get_magic(new_code)
268
269        # Update kernel configuration first, if needed.
270        config_magic = new_magic.get("config")
271        if config_magic is not None:
272            self.parse_config_magic(config_magic)
273
274        if self.should_reset(new_magic):
275            self._previous_code = new_code
276            self._previous_magic = new_magic
277        else:
278            self._previous_code += ("\n" if self._previous_code else "") + new_code
279            self._previous_magic.update(new_magic)
280
281        return self._previous_code, self._previous_magic.get("args", [])
282
283    def make_status(self):
284        return {
285            "status": "ok",
286            "execution_count": self.execution_count,
287            "payload": [],
288            "user_expressions": {},
289        }
290
291    def send_stream(self, name, content):
292        self.send_response(self.iopub_socket, "stream", {"name": name, "text": content})
293
294        return self.make_status()
295
296    def send_stderr(self, stderr):
297        return self.send_stream("stderr", stderr)
298
299    def send_stdout(self, stdout):
300        return self.send_stream("stdout", stdout)
301
302    def do_execute(
303        self, code, silent, store_history=True, user_expressions=None, allow_stdin=False
304    ):
305        """Execute user code using llvm-tblgen binary."""
306        try:
307            all_code, args = self.get_code_and_args(code)
308        except TableGenKernelException as e:
309            return self.send_stderr(str(e))
310
311        # If we cannot find llvm-tblgen, propagate the error to the notebook.
312        # (in case the user is not able to see the output from the Jupyter server)
313        try:
314            executable = self.get_executable()
315        except Exception as e:
316            return self.send_stderr(str(e))
317
318        with tempfile.TemporaryFile("w+") as f:
319            f.write(all_code)
320            f.seek(0)
321            got = subprocess.run(
322                [executable, *args],
323                stdin=f,
324                stderr=subprocess.PIPE,
325                stdout=subprocess.PIPE,
326                universal_newlines=True,
327            )
328
329        if not silent:
330            if got.stderr:
331                return self.send_stderr(got.stderr)
332            else:
333                return self.send_stdout(got.stdout)
334        else:
335            return self.make_status()
336
337
338if __name__ == "__main__":
339    import doctest
340
341    doctest.testmod(extraglobs={"k": TableGenKernel()})
342