-
-
Notifications
You must be signed in to change notification settings - Fork 93
/
ob-jupyter.el
836 lines (732 loc) · 36.7 KB
/
ob-jupyter.el
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
;;; ob-jupyter.el --- Jupyter integration with org-mode -*- lexical-binding: t -*-
;; Copyright (C) 2018-2024 Nathaniel Nicandro
;; Author: Nathaniel Nicandro <[email protected]>
;; Created: 21 Jan 2018
;; This program is free software; you can redistribute it and/or
;; modify it under the terms of the GNU General Public License as
;; published by the Free Software Foundation; either version 3, or (at
;; your option) any later version.
;; This program is distributed in the hope that it will be useful, but
;; WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
;; General Public License for more details.
;; You should have received a copy of the GNU General Public License
;; along with GNU Emacs; see the file COPYING. If not, write to the
;; Free Software Foundation, Inc., 59 Temple Place - Suite 330,
;; Boston, MA 02111-1307, USA.
;;; Commentary:
;; Interact with a Jupyter kernel via `org-mode' src-block's.
;;; Code:
(defgroup ob-jupyter nil
"Jupyter integration with org-mode"
:group 'org-babel)
(require 'jupyter-env)
(require 'jupyter-kernelspec)
(require 'jupyter-org-client)
(require 'jupyter-org-extensions)
(eval-when-compile
(require 'jupyter-repl) ; For `jupyter-with-repl-buffer'
(require 'subr-x))
(declare-function org-in-src-block-p "org" (&optional inside))
(declare-function org-element-at-point "org-element")
(declare-function org-element-property "org-element" (property element))
(declare-function org-element-context "org-element" (&optional element))
(declare-function org-babel-execute-src-block "ob-core" (&optional arg info params executor-type))
(declare-function org-babel-variable-assignments:python "ob-python" (params))
(declare-function org-babel-expand-body:generic "ob-core" (body params &optional var-lines))
(declare-function org-export-derived-backend-p "ox" (backend &rest backends))
(declare-function jupyter-run-server-repl "jupyter-server")
(declare-function jupyter-connect-server-repl "jupyter-server")
(declare-function jupyter-kernelspecs "jupyter-server")
(declare-function jupyter-server-kernel-id-from-name "jupyter-server")
(declare-function jupyter-server-name-client-kernel "jupyter-server")
(declare-function jupyter-api-get-kernel "jupyter-rest-api")
(declare-function jupyter-tramp-url-from-file-name "jupyter-tramp")
(declare-function jupyter-tramp-server-from-file-name "jupyter-tramp")
(declare-function jupyter-tramp-file-name-p "jupyter-tramp")
(defcustom org-babel-jupyter-language-aliases '(("python3" "python"))
"An alist mapping kernel language names to another name.
If a kernel has a language name matching the CAR of an element of
this list, the associated name will be used for the names of the
source blocks instead.
So if this variable has an entry like \\='(\"python3\" \"python\")
then instead of jupyter-python3 source blocks, you can use
jupyter-python source blocks for the associated kernel."
:type '(alist :key-type string :value-type string))
(defvaralias 'org-babel-jupyter-resource-directory
'jupyter-org-resource-directory)
(defvar org-babel-jupyter-session-clients (make-hash-table :test #'equal)
"A hash table mapping session names to Jupyter clients.
`org-babel-jupyter-src-block-session' returns a key into this
table for the source block at `point'.")
(defvar org-babel-header-args:jupyter '((kernel . :any)
(async . ((yes no))))
"Available header arguments for Jupyter src-blocks.")
(defvar org-babel-default-header-args:jupyter '((:kernel . "python")
(:async . "no"))
"Default header arguments for Jupyter src-blocks.")
;;; Helper functions
(defun org-babel-jupyter--src-block-kernel-language ()
(when (org-in-src-block-p)
(let ((info (org-babel-get-src-block-info)))
(save-match-data
(string-match "^jupyter-\\(.+\\)$" (car info))
(match-string 1 (car info))))))
(defun org-babel-jupyter-language-p (lang)
"Return non-nil if LANG src-blocks are executed using Jupyter."
(or (string-prefix-p "jupyter-" lang)
;; Check if the language has been overridden, see
;; `org-babel-jupyter-override-src-block'
(advice-member-p
'ob-jupyter (intern (concat "org-babel-execute:" lang)))))
(defun org-babel-jupyter-session-key (params)
"Return a string that is the concatenation of the :session and :kernel PARAMS.
PARAMS is the arguments alist as returned by
`org-babel-get-src-block-info'. The returned string can then be
used to identify unique Jupyter Org babel sessions."
;; Take into account a Lisp expression as a session name.
(let ((session (org-babel-read (alist-get :session params)))
(kernel (alist-get :kernel params)))
(unless (and session kernel
(not (equal session "none")))
(error "Need a valid session and a kernel to form a key"))
(concat session "-" kernel)))
(defun org-babel-jupyter-src-block-session ()
"Return the session key for the current Jupyter source block.
Return nil if the current source block is not a Jupyter block or
if there is no source block at point."
(let ((info (or (and (org-in-src-block-p)
(org-babel-get-src-block-info 'light))
(org-babel-lob-get-info))))
(when info
(cl-destructuring-bind (lang _ params . rest) info
(when (org-babel-jupyter-language-p lang)
(org-babel-jupyter-session-key params))))))
;;; `ob' integration
(defun org-babel-variable-assignments:jupyter (params &optional lang)
"Assign variables in PARAMS according to the Jupyter kernel language.
LANG is the kernel language of the source block. If LANG is nil,
get the kernel language from the current source block.
The variables are assigned by looking for the function
`org-babel-variable-assignments:LANG'. If this function does not
exist or if LANG cannot be determined, assign variables using
`org-babel-variable-assignments:python'."
(or lang (setq lang (org-babel-jupyter--src-block-kernel-language)))
(let ((fun (when lang
(intern (format "org-babel-variable-assignments:%s" lang)))))
(if (functionp fun) (funcall fun params)
(require 'ob-python)
(org-babel-variable-assignments:python params))))
(cl-defgeneric org-babel-jupyter-transform-code (code _changelist)
"Transform CODE according to CHANGELIST, return the transformed CODE.
CHANGELIST is a property list containing the requested changes. The default
implementation returns CODE unchanged.
This is useful for kernel languages to extend using the
jupyter-lang method specializer, e.g. to return new code to change
directories before evaluating CODE.
See `org-babel-expand-body:jupyter' for possible changes that can
be in CHANGELIST."
code)
(defun org-babel-expand-body:jupyter (body params &optional var-lines lang)
"Expand BODY according to PARAMS.
BODY is the code to expand, PARAMS should be the header arguments
of the src block with BODY as its code, and VAR-LINES should be
the list of strings containing the variables to evaluate before
executing body. LANG is the kernel language of the source block.
This function is similar to
`org-babel-variable-assignments:jupyter' in that it attempts to
find the kernel language of the source block if LANG is not
provided.
BODY is expanded by calling the function
`org-babel-expand-body:LANG'. If this function doesn't exist or
if LANG cannot be determined, fall back to
`org-babel-expand-body:generic'.
If PARAMS has a :dir parameter, the expanded code is passed to
`org-babel-jupyter-transform-code' with a changelist that
includes the :dir parameter with the directory being an absolute
path."
(or lang (setq lang (org-babel-jupyter--src-block-kernel-language)))
(let* ((expander (when lang
(intern (format "org-babel-expand-body:%s" lang))))
(expanded (if (functionp expander)
(funcall expander body params)
(org-babel-expand-body:generic body params var-lines)))
(changelist nil))
(when-let* ((dir (alist-get :dir params)))
(setq changelist (plist-put changelist :dir (expand-file-name dir))))
(if changelist (org-babel-jupyter-transform-code expanded changelist)
expanded)))
(defun org-babel-edit-prep:jupyter (info)
"Prepare the edit buffer according to INFO.
Enable `jupyter-repl-interaction-mode' in the edit buffer
associated with the session found in INFO.
If the session is a Jupyter TRAMP file name, the
`default-directory' of the edit buffer is set to the root
directory the notebook serves.
If `jupyter-org-auto-connect' is nil, this function does nothing
if the session has not been initiated yet."
(let* ((params (nth 2 info))
(session (alist-get :session params))
(client-buffer
(when (or jupyter-org-auto-connect
(org-babel-jupyter-session-initiated-p params))
(org-babel-jupyter-initiate-session session params))))
(when client-buffer
(jupyter-repl-associate-buffer client-buffer)
(when (jupyter-tramp-file-name-p session)
(setq default-directory (concat (file-remote-p session) "/"))))))
(defun org-babel-jupyter--insert-variable-assignments (params)
"Insert variable assignment lines from PARAMS into the `current-buffer'.
Return non-nil if there are variable assignments, otherwise
return nil."
(let ((var-lines (org-babel-variable-assignments:jupyter params)))
(prog1 var-lines
(jupyter-repl-replace-cell-code (mapconcat #'identity var-lines "\n")))))
(defun org-babel-prep-session:jupyter (session params)
"Prepare a Jupyter SESSION according to PARAMS."
(with-current-buffer (org-babel-jupyter-initiate-session session params)
(goto-char (point-max))
(and (org-babel-jupyter--insert-variable-assignments params)
(jupyter-repl-execute-cell jupyter-current-client))
(current-buffer)))
(defun org-babel-load-session:jupyter (session body params)
"In a Jupyter SESSION, load BODY according to PARAMS."
(save-window-excursion
(with-current-buffer (org-babel-jupyter-initiate-session session params)
(goto-char (point-max))
(when (org-babel-jupyter--insert-variable-assignments params)
(insert "\n"))
(insert (org-babel-expand-body:jupyter (org-babel-chomp body) params))
(current-buffer))))
(defvar org-babel-jupyter-resolving-reference-p nil
"Non-nil if a reference is being resolved.")
(defun org-babel-jupyter--indicate-resolve (&rest args)
"Set `org-babel-jupyter-resolving-referece-p', apply ARGS."
(let ((org-babel-jupyter-resolving-reference-p t))
(apply args)))
(advice-add #'org-babel-ref-resolve :around #'org-babel-jupyter--indicate-resolve)
;;;; Initializing session clients
(cl-defstruct (org-babel-jupyter-session
(:constructor org-babel-jupyter-session))
name)
(cl-defstruct (org-babel-jupyter-remote-session
(:include org-babel-jupyter-session)
(:constructor org-babel-jupyter-remote-session))
connect-repl-p)
(cl-defmethod org-babel-jupyter-parse-session ((session string))
"Return a parsed representation of SESSION."
(org-babel-jupyter-session :name session))
(cl-defmethod org-babel-jupyter-initiate-client ((_session org-babel-jupyter-session) kernel)
"Launch SESSION's KERNEL, return a `jupyter-org-client' connected to it.
SESSION is the :session header argument of a source block and
KERNEL is the name of the kernel to launch."
(jupyter-run-repl kernel nil nil 'jupyter-org-client))
(cl-defmethod org-babel-jupyter-initiate-client :around (session _kernel)
"Rename the returned client's REPL buffer to include SESSION's name.
Also set `jupyter-include-other-output' to nil for the session so
that output produced by other clients do not get handled by the
client."
(let ((client (cl-call-next-method)))
(prog1 client
(jupyter-set client 'jupyter-include-other-output nil)
;; Append the name of SESSION to the initiated client REPL's
;; `buffer-name'.
(jupyter-with-repl-buffer client
(let ((name (buffer-name)))
(when (string-match "^\\*\\(.+\\)\\*" name)
(rename-buffer
(concat "*" (match-string 1 name) "-"
(org-babel-jupyter-session-name session)
"*")
'unique)))))))
(cl-defmethod org-babel-jupyter-parse-session :extra "remote" ((session string))
"If SESSION is a remote file name, return a `org-babel-jupyter-remote-session'.
A `org-babel-jupyter-remote-session' is also returned if SESSION
ends in \".json\", regardless of SESSION being a remote file
name, with `org-babel-jupyter-remote-session-connect-repl-p' set
to nil. The CONNECT-REPL-P slot indicates that a connection file
is read to connect to the session, as opposed to launching a
kernel."
(if jupyter-use-zmq
(let ((json-p (string-suffix-p ".json" session)))
(if (or json-p (file-remote-p session))
(org-babel-jupyter-remote-session
:name session
:connect-repl-p json-p)
(cl-call-next-method)))
(when (file-remote-p session)
(error "ZMQ is required for remote sessions (%s)" session))
(cl-call-next-method)))
(cl-defmethod org-babel-jupyter-initiate-client :before ((session org-babel-jupyter-remote-session) _kernel)
"Raise an error if SESSION's name is a remote file name without a local name.
The local name is used as a unique identifier of a remote
session."
(unless (not (zerop (length (file-local-name
(org-babel-jupyter-session-name session)))))
(error "No remote session name")))
(cl-defmethod org-babel-jupyter-initiate-client ((session org-babel-jupyter-remote-session) kernel)
"Initiate a client connected to a remote kernel process."
(pcase-let (((cl-struct org-babel-jupyter-remote-session name connect-repl-p) session))
(if connect-repl-p
(jupyter-connect-repl name nil nil 'jupyter-org-client)
(let ((default-directory (file-remote-p name)))
(org-babel-jupyter-aliases-from-kernelspecs)
(jupyter-run-repl kernel nil nil 'jupyter-org-client)))))
(require 'jupyter-server)
(require 'jupyter-tramp)
(cl-defstruct (org-babel-jupyter-server-session
(:include org-babel-jupyter-remote-session)
(:constructor org-babel-jupyter-server-session)))
(cl-defmethod org-babel-jupyter-parse-session :extra "server" ((session string))
"If SESSION is a Jupyter TRAMP file name return a
`org-babel-jupyter-server-session'."
(if (jupyter-tramp-file-name-p session)
(org-babel-jupyter-server-session :name session)
(cl-call-next-method)))
(cl-defmethod org-babel-jupyter-initiate-client ((session org-babel-jupyter-server-session) kernel)
(let* ((rsession (org-babel-jupyter-session-name session))
(server (with-parsed-tramp-file-name rsession nil
(when (member host '("127.0.0.1" "localhost"))
(setq port (tramp-file-name-port-or-default v))
(when (jupyter-port-available-p port)
(if (y-or-n-p (format "Notebook not started on port %s. Launch one? "
port))
;; TODO: Specify authentication? But then
;; how would you get the token for the
;; login that happens in
;; `jupyter-tramp-server-from-file-name'.
(jupyter-launch-notebook port)
(user-error "Launch a notebook on port %s first." port))))
(jupyter-tramp-server-from-file-name rsession))))
(unless (jupyter-server-has-kernelspec-p server kernel)
(error "No kernelspec matching \"%s\" exists at %s"
kernel (oref server url)))
;; Language aliases may not exist for the kernels that are accessible on
;; the server so ensure they do.
(org-babel-jupyter-aliases-from-kernelspecs
nil (jupyter-kernelspecs server))
(let ((sname (file-local-name rsession)))
(if-let ((id (jupyter-server-kernel-id-from-name server sname)))
;; Connecting to an existing kernel
(cl-destructuring-bind (&key name id &allow-other-keys)
(or (ignore-errors (jupyter-api-get-kernel server id))
(error "Kernel ID, %s, no longer references a kernel at %s"
id (oref server url)))
(unless (string-match-p kernel name)
(error "\":kernel %s\" doesn't match \"%s\"" kernel name))
(jupyter-connect-server-repl server id nil nil 'jupyter-org-client))
;; Start a new kernel
(let ((client (jupyter-run-server-repl
server kernel nil nil 'jupyter-org-client)))
(prog1 client
;; TODO: If a kernel gets renamed in the future it doesn't affect
;; any source block :session associations because the hash of the
;; session name used here is already stored in the
;; `org-babel-jupyter-session-clients' variable. Should that
;; variable be updated on a kernel rename?
;;
;; TODO: Would we always want to do this?
(jupyter-server-name-client-kernel client sname)))))))
(defun org-babel-jupyter-session-initiated-p (params)
"Return non-nil if the session corresponding to PARAMS is initiated."
(let ((key (org-babel-jupyter-session-key params)))
(gethash key org-babel-jupyter-session-clients)))
(defun org-babel-jupyter-initiate-session-by-key (session params)
"Return the Jupyter REPL buffer for SESSION.
If SESSION does not have a client already, one is created based
on SESSION and PARAMS. If SESSION ends with \".json\" then
SESSION is interpreted as a kernel connection file and a new
kernel connected to SESSION is created.
Otherwise a kernel is started based on the `:kernel' parameter in
PARAMS which should be either a valid kernel name or a prefix of
one, in which case the first kernel that matches the prefix will
be used.
If SESSION is a remote file name, like /ssh:ec2:jl, then the
kernel starts on the remote host /ssh:ec2: with a session name of
jl. The remote host must have jupyter installed since the
\"jupyter kernel\" command will be used to start the kernel on
the host."
(let* ((key (org-babel-jupyter-session-key params))
(client (gethash key org-babel-jupyter-session-clients)))
(unless client
(setq client (org-babel-jupyter-initiate-client
(org-babel-jupyter-parse-session session)
(alist-get :kernel params)))
(puthash key client org-babel-jupyter-session-clients)
(jupyter-with-repl-buffer client
(let ((forget-client (lambda () (remhash key org-babel-jupyter-session-clients))))
(add-hook 'kill-buffer-hook forget-client nil t))))
(oref client buffer)))
(defun org-babel-jupyter-initiate-session (&optional session params)
"Initialize a Jupyter SESSION according to PARAMS."
(if (equal session "none") (error "Need a session to run")
(when session
;; Take into account a Lisp expression as a session name.
(setq session (org-babel-read session)))
(org-babel-jupyter-initiate-session-by-key session params)))
;;;; Helper functions
;;;###autoload
(defun org-babel-jupyter-scratch-buffer ()
"Display a scratch buffer connected to the current block's session."
(interactive)
(let (buffer)
(org-babel-do-in-edit-buffer
(setq buffer (save-window-excursion
(jupyter-repl-scratch-buffer))))
(if buffer (pop-to-buffer buffer)
(user-error "No source block at point"))))
(cl-defmethod jupyter-do-refresh-kernelspecs (&context (major-mode org-mode))
(or (jupyter-org-when-in-src-block
(let* ((info (org-babel-get-src-block-info 'light))
(params (nth 2 info))
(session (org-babel-read (alist-get :session params))))
(when (file-remote-p session)
(jupyter-kernelspecs session 'refresh))))
(cl-call-next-method)))
;;;; `org-babel-execute:jupyter'
(defvar org-link-bracket-re)
(defun org-babel-jupyter-cleanup-file-links ()
"Delete the files of image links for the current source block result.
Do this only if the file exists in
`org-babel-jupyter-resource-directory'."
(when-let*
((pos (org-babel-where-is-src-block-result))
(link-re (format "^[ \t]*%s[ \t]*$" org-link-bracket-re))
(resource-dir (expand-file-name org-babel-jupyter-resource-directory)))
(save-excursion
(goto-char pos)
(forward-line)
(let ((bound (org-babel-result-end)))
;; This assumes that `jupyter-org-client' only emits bracketed links as
;; images
(while (re-search-forward link-re bound t)
(when-let*
((path (org-element-property :path (org-element-context)))
(dir (when (file-name-directory path)
(expand-file-name (file-name-directory path)))))
(when (and (equal dir resource-dir)
(file-exists-p path))
(delete-file path))))))))
;; TODO: What is a better way to handle discrepancies between how `org-mode'
;; views header arguments and how `emacs-jupyter' views them? Should the
;; strategy be to always try to emulate the `org-mode' behavior?
(defun org-babel-jupyter--remove-file-param (params)
"Destructively remove the file result parameter from PARAMS.
These parameters are handled internally."
(let* ((result-params (assq :result-params params))
(fresult (member "file" result-params))
(fparam (assq :file params)))
(setcar fresult "")
(delq fparam params)))
(defconst org-babel-jupyter-async-inline-results-pending-indicator "???"
"A string to disambiguate pending inline results from empty results.")
(defun org-babel-jupyter--execute (code async-p)
(jupyter-run-with-client jupyter-current-client
(let ((dreq (jupyter-execute-request :code code)))
(jupyter-mlet* ((req (jupyter-org-maybe-queued dreq)))
(jupyter-return
`(,req
,(cond
(async-p
(when (bound-and-true-p org-export-current-backend)
(jupyter-add-idle-sync-hook
'org-babel-after-execute-hook req 'append))
(if (jupyter-org-request-inline-block-p req)
org-babel-jupyter-async-inline-results-pending-indicator
;; This returns the message ID of REQ as an indicator
;; for the pending results.
(jupyter-org-pending-async-results req)))
(t
(jupyter-idle-sync req)
(if (jupyter-org-request-inline-block-p req)
;; When evaluating a source block synchronously, only the
;; :execute-result will be in `jupyter-org-request-results' since
;; stream results and any displayed data will be placed in a separate
;; buffer.
(let ((el (jupyter-org-result
req (car (jupyter-org-request-results req)))))
(if (stringp el) el
(org-element-property :value el)))
;; This returns an Org formatted string of the collected
;; results.
(jupyter-org-sync-results req))))))))))
(defvar org-babel-jupyter-current-src-block-params nil
"The block parameters of the most recently executed Jupyter source block.")
(defun org-babel-execute:jupyter (body params)
"Execute BODY according to PARAMS.
BODY is the code to execute for the current Jupyter `:session' in
the PARAMS alist."
(when org-babel-current-src-block-location
(save-excursion
(goto-char org-babel-current-src-block-location)
(when (jupyter-org-request-at-point)
(user-error "Source block currently being executed"))))
(let* ((result-params (assq :result-params params))
(async-p (jupyter-org-execute-async-p params)))
(when (member "replace" result-params)
(org-babel-jupyter-cleanup-file-links))
(let* ((org-babel-jupyter-current-src-block-params params)
(session (alist-get :session params))
(buf (org-babel-jupyter-initiate-session session params))
(jupyter-current-client (buffer-local-value 'jupyter-current-client buf))
(lang (jupyter-kernel-language jupyter-current-client))
(vars (org-babel-variable-assignments:jupyter params lang))
(code (progn
(when-let* ((dir (alist-get :dir params)))
;; `default-directory' is already set according
;; to :dir when executing a source block. Set
;; :dir to the absolute path so that
;; `org-babel-expand-body:jupyter' does not try
;; to re-expand the path. See #302.
(setf (alist-get :dir params) default-directory))
(org-babel-expand-body:jupyter body params vars lang))))
(pcase-let ((`(,req ,maybe-result)
(org-babel-jupyter--execute code async-p)))
;; KLUDGE: Remove the file result-parameter so that
;; `org-babel-insert-result' doesn't attempt to handle it while
;; async results are pending. Do the same in the synchronous
;; case, but not if link or graphics are also result-parameters,
;; only in Org >= 9.2, since those in combination with file mean
;; to interpret the result as a file link, a useful meaning that
;; doesn't interfere with Jupyter style result insertion.
;;
;; Do this after sending the request since
;; `jupyter-generate-request' still needs access to the :file
;; parameter.
(when (and (member "file" result-params)
(or async-p
(not (or (member "link" result-params)
(member "graphics" result-params)))))
(org-babel-jupyter--remove-file-param params))
(prog1 maybe-result
;; KLUDGE: Add the "raw" result parameter for non-inline
;; synchronous results because an Org formatted string is
;; already returned in that case and
;; `org-babel-insert-result' should not process it.
(unless (or async-p
(jupyter-org-request-inline-block-p req))
(nconc (alist-get :result-params params) (list "raw"))))))))
;;; Overriding source block languages, language aliases
(defvar org-babel-jupyter--babel-ops
'(execute expand-body prep-session edit-prep
variable-assignments load-session
initiate))
(defvar org-babel-jupyter--babel-vars
'(header-args default-header-args))
(defun org-babel-jupyter--babel-op-symbol (op lang)
(if (eq op 'initiate)
(intern (format "org-babel-%s-initiate-session" lang))
(intern (format (format "org-babel-%s:%s" op lang)))))
(defun org-babel-jupyter--babel-var-symbol (var lang)
(intern (format "org-babel-%s:%s" var lang)))
(defun org-babel-jupyter--babel-map (alias-action
var-action)
"Loop over Org babel function and variable symbols.
ALIAS-ACTION and VAR-ACTION are functions of one argument.
When ALIAS-ACTION is called, the argument will be a symbol that
represents an Org Babel operation that can be defined by a
language extension to Org Babel, e.g. \\='execute.
Similarly VAR-ACTION is called with a symbol representing an Org
Babel variable that can be defined for a language,
e.g. \\='header-args."
(declare (indent 0))
(dolist (op org-babel-jupyter--babel-ops)
(funcall alias-action op))
(dolist (var org-babel-jupyter--babel-vars)
(funcall var-action var)))
(defun org-babel-jupyter-override-src-block (lang)
"Override the built-in `org-babel' functions for LANG.
This overrides functions like `org-babel-execute:LANG' and
`org-babel-LANG-initiate-session' to use the machinery of
jupyter-LANG source blocks.
Also, set `org-babel-header-args:LANG' to the value of
`org-babel-header-args:jupyter-LANG', if the latter exists. If
`org-babel-header-args:LANG' had a value, save it as a symbol
property of `org-babel-header-args:LANG' for restoring it later.
Do the same for `org-babel-default-header-args:LANG'."
(org-babel-jupyter--babel-map
(lambda (op)
;; Only override operations that are not related to a particular
;; language.
(unless (memq op '(variable-assignments expand-body))
(let ((lang-op
(org-babel-jupyter--babel-op-symbol
op lang))
(jupyter-lang-op
(org-babel-jupyter--babel-op-symbol
op (format "jupyter-%s" lang))))
;; If a language doesn't have a function assigned, set one so it can
;; be overridden
(unless (fboundp lang-op)
(fset lang-op #'ignore))
(advice-add lang-op :override jupyter-lang-op
'((name . ob-jupyter))))))
(lambda (var)
(let ((lang-var
(org-babel-jupyter--babel-var-symbol
var lang))
(jupyter-lang-var
(org-babel-jupyter--babel-var-symbol
var (format "jupyter-%s" lang))))
(when (boundp jupyter-lang-var)
(when (boundp lang-var)
(put lang-var 'jupyter-restore-value (symbol-value lang-var)))
(set lang-var (copy-tree (symbol-value jupyter-lang-var))))))))
(defun org-babel-jupyter-restore-src-block (lang)
"Restore the overridden `org-babel' functions for LANG.
This undoes everything that
`org-babel-jupyter-override-src-block' did."
(org-babel-jupyter--babel-map
(lambda (op)
;; Only override operations that are not related to a particular
;; language.
(unless (memq op '(variable-assignments expand-body))
(let ((lang-op
(org-babel-jupyter--babel-op-symbol
op lang))
(jupyter-lang-op
(org-babel-jupyter--babel-op-symbol
op (format "jupyter-%s" lang))))
(advice-remove lang-op jupyter-lang-op)
;; The function didn't have a definition, so
;; ensure that we restore that fact.
(when (eq (symbol-function lang-op) #'ignore)
(fmakunbound lang-op)))))
(lambda (var)
(let ((lang-var
(org-babel-jupyter--babel-var-symbol
var lang)))
(when (boundp lang-var)
(set lang-var (get lang-var 'jupyter-restore-value)))))))
(defun org-babel-jupyter-make-language-alias (kernel lang)
"Similar to `org-babel-make-language-alias' but for Jupyter src-blocks.
KERNEL should be the name of the default kernel to use for kernel
LANG, the language of the kernel.
The Org Babel functions `org-babel-FN:jupyter-LANG', where FN is
one of execute, expand-body, prep-session, edit-prep,
variable-assignments, or load-session, are aliased to
`org-babel-FN:jupyter'. Similarly,
`org-babel-jupyter-LANG-initiate-session' is aliased to
`org-babel-jupyter-initiate-session'.
If not already defined, the variable
`org-babel-default-header-args:jupyter-LANG' is set to the same
value as `org-babel-header-args:jupyter', which see. The
variable `org-babel-default-header-args:jupyter-LANG' is also set
to
\((:async . \"no\")
\(:kernel . KERNEL))
if that variable does not already have a value.
If LANG has an association in `org-babel-tangle-lang-exts',
associate the same value with jupyter-LANG, if needed.
Similarly, associate the same value for LANG in
`org-src-lang-modes'."
(org-babel-jupyter--babel-map
(lambda (op)
(defalias (org-babel-jupyter--babel-op-symbol
op (format "jupyter-%s" lang))
(org-babel-jupyter--babel-op-symbol
op "jupyter")))
(lambda (var)
(let ((jupyter-var
(org-babel-jupyter--babel-var-symbol
var "jupyter"))
(jupyter-lang-var
(org-babel-jupyter--babel-var-symbol
var (format "jupyter-%s" lang))))
(unless (boundp jupyter-lang-var)
(set jupyter-lang-var (copy-tree (symbol-value jupyter-var)))
(cond
((eq var 'default-header-args)
;; Needed since the default kernel is not language
;; specific and it needs to be.
(setf (alist-get :kernel (symbol-value jupyter-lang-var)) kernel)
(put jupyter-lang-var 'variable-documentation
(format
"Default header arguments for Jupyter %s src-blocks"
lang)))
(t
(put jupyter-lang-var 'variable-documentation
(get jupyter-var 'variable-documentation))))))))
(when (assoc lang org-babel-tangle-lang-exts)
(add-to-list 'org-babel-tangle-lang-exts
(cons (concat "jupyter-" lang)
(cdr (assoc lang org-babel-tangle-lang-exts)))))
(add-to-list 'org-src-lang-modes
(cons (concat "jupyter-" lang)
(or (cdr (assoc lang org-src-lang-modes))
(intern (downcase (replace-regexp-in-string
"[0-9]*" "" lang)))))))
(defun org-babel-jupyter-aliases-from-kernelspecs (&optional refresh specs)
"Make language aliases based on the available kernelspecs.
For all kernel SPECS, make a language alias for the kernel
language if one does not already exist. The alias is created with
`org-babel-jupyter-make-language-alias'.
SPECS defaults to those associated with the `default-directory'.
Optional argument REFRESH has the same meaning as in
`jupyter-kernelspecs'.
Note, spaces in the kernel language name are converted into
dashes in the language alias, e.g.
Wolfram Language -> jupyter-Wolfram-Language
For convenience, after creating a language alias for a kernel
language LANG, set the :kernel default header argument if not
present in `org-babel-default-header-args:jupyter-LANG', see
`org-babel-header-args:jupyter'. This allows users to set that
variable in their configurations without having to also set the
:kernel header argument since it is common for only one per
language to exist on someone's system."
(cl-loop
for spec in (or specs
(with-demoted-errors "Error retrieving kernelspecs: %S"
(jupyter-kernelspecs default-directory refresh)))
for kernel = (jupyter-kernelspec-name spec)
for lang = (let ((lang (jupyter-canonicalize-language-string
(plist-get (jupyter-kernelspec-plist spec) :language))))
(or (cadr (assoc lang org-babel-jupyter-language-aliases))
lang))
unless (member lang languages) collect lang into languages and
do (org-babel-jupyter-make-language-alias kernel lang)
;; KLUDGE: The :kernel header argument is always set, even when we aren't
;; the ones who originally set the defaults. This is here for convenience
;; since usually a user does not set :kernel directly.
(let ((var (intern (concat "org-babel-default-header-args:jupyter-" lang))))
(unless (alist-get :kernel (symbol-value var))
(setf (alist-get :kernel (symbol-value var)) kernel)))))
;;; `ox' integration
(defvar org-latex-minted-langs)
(defun org-babel-jupyter-setup-export (backend)
"Ensure that Jupyter src-blocks are integrated with BACKEND.
Currently this makes sure that Jupyter src-block languages are
mapped to their appropriate minted language in
`org-latex-minted-langs' if BACKEND is latex."
(cond
((org-export-derived-backend-p backend 'latex)
(cl-loop
for spec in (jupyter-kernelspecs default-directory)
for lang = (plist-get (jupyter-kernelspec-plist spec) :language)
do (cl-pushnew (list (intern (concat "jupyter-" lang)) lang)
org-latex-minted-langs :test #'equal)))))
(defun org-babel-jupyter-strip-ansi-escapes (_backend)
"Remove ANSI escapes from Jupyter src-block results in the current buffer."
(org-babel-map-src-blocks nil
(when (org-babel-jupyter-language-p lang)
(when-let* ((pos (org-babel-where-is-src-block-result))
(ansi-color-apply-face-function
(lambda (beg end face)
;; Could be useful for export backends
(when face
(put-text-property beg end 'face face)))))
(goto-char pos)
(ansi-color-apply-on-region (point) (org-babel-result-end))))))
;;; Hook into `org'
;; Defer generation of the aliases until Org is enabled in a buffer to
;; avoid generating them at top-level when loading ob-jupyter. Some
;; users, e.g. those who use conda environments, may not have a
;; jupyter command available at load time.
(defun org-babel-jupyter-make-local-aliases ()
(let ((default-directory user-emacs-directory))
(org-babel-jupyter-aliases-from-kernelspecs)))
(add-hook 'org-mode-hook #'org-babel-jupyter-make-local-aliases 10)
(add-hook 'org-export-before-processing-functions #'org-babel-jupyter-setup-export)
(add-hook 'org-export-before-parsing-functions #'org-babel-jupyter-strip-ansi-escapes)
(provide 'ob-jupyter)
;;; ob-jupyter.el ends here