From 580776bf58be1e4605c3a650643fb9771f295b8f Mon Sep 17 00:00:00 2001
From: Nathaniel Nicandro <nathanielnicandro@gmail.com>
Date: Sat, 23 Mar 2019 00:12:06 -0500
Subject: [PATCH] Add `org-babel-jupyter-src-block-session`

---
 ob-jupyter.el        | 25 ++++++++++++++++++++++++-
 test/jupyter-test.el | 24 ++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/ob-jupyter.el b/ob-jupyter.el
index 3a9ca20..f3b83f4 100644
--- a/ob-jupyter.el
+++ b/ob-jupyter.el
@@ -78,6 +78,29 @@
       (advice-member-p
        'ob-jupyter (intern (concat "org-babel-execute:" lang)))))
 
+(defun org-babel-jupyter-session-key (params)
+  "Return the session key based on the keys in PARAMS.
+PARAMS is the arguments alist as returned by
+`org-babel-get-src-block-info' and should contain a :kernel key
+and a valid :session key. The session key is used to access the
+clients in `org-babel-jupyter-session-clients'."
+  (let ((session (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."
+  (when (org-in-src-block-p)
+    (cl-destructuring-bind (lang _ params . rest)
+        (org-babel-get-src-block-info 'light)
+      (when (org-babel-jupyter-language-p lang)
+        (org-babel-jupyter-session-key params)))))
+
 ;;; `ob' integration
 
 (defun org-babel-variable-assignments:jupyter (params &optional lang)
@@ -189,7 +212,7 @@ of one. The first kernel that is returned by
 `jupyter-find-kernelspecs' when passed the value of the `:kernel'
 parameter will be used."
   (let* ((kernel (alist-get :kernel params))
-         (key (concat session "-" kernel))
+         (key (org-babel-jupyter-session-key params))
          (client
           (or (gethash key org-babel-jupyter-session-clients)
               (let ((client
diff --git a/test/jupyter-test.el b/test/jupyter-test.el
index da2332d..f880912 100644
--- a/test/jupyter-test.el
+++ b/test/jupyter-test.el
@@ -1712,6 +1712,30 @@ ABCDEF
       (remove-text-properties (point-min) (point-max) '(face))
       (funcall test-fun 24 '(19 20 21 22 23 26 27 28 29)))))
 
+(ert-deftest org-babel-jupyter-src-block-session ()
+  :tags '(org)
+  (jupyter-org-test
+   (insert "\
+#+BEGIN_SRC jupyter-foo :session bar
+#+END_SRC")
+   (goto-char (point-min))
+   (should-error (org-babel-jupyter-src-block-session))
+   (erase-buffer)
+   (insert "\
+#+BEGIN_SRC jupyter-foo :kernel bar
+#+END_SRC")
+   (goto-char (point-min))
+   (should-error (org-babel-jupyter-src-block-session))
+   (erase-buffer)
+   (insert "\
+#+BEGIN_SRC jupyter-foo :session bar :kernel bar
+#+END_SRC")
+   (goto-char (point-min))
+   (should (equal (org-babel-jupyter-src-block-session)
+                  (org-babel-jupyter-session-key
+                   (nth 2 (org-babel-get-src-block-info 'light)))))
+   (erase-buffer)))
+
 (ert-deftest org-babel-jupyter-:results-header-arg ()
   :tags '(org)
   (ert-info ("scalar suppresses table output")