scimax-jupyter.el

;;; scimax-jupyter.el --- scimax customization for emacs-jupyter

;;; Commentary:
;; This provides some fine-tuning of key bindings, and loads jupyter for scimax.

(require 'scimax-ob)
(require 'jupyter)

;; * automatic buffer kernels This is done by setting a session name in the org
;; file <jps is a yasnippet that does this reliably. I decided for now it is too
;; difficult to advise everything in jupyter to get the same result.


;; * scimax jupyter header defaults

(add-to-list 'org-babel-load-languages '(jupyter . t) t)

(setq org-babel-default-header-args:jupyter-python
      '((:results . "both")
	(:session . "jupyter-python")
	(:kernel . "python3")
	(:pandoc . "t")
	(:exports . "both")
	(:cache .   "no")
	(:noweb . "no")
	(:hlines . "no")
	(:tangle . "no")
	(:eval . "never-export")))


;; ** Julia
;; (setq org-babel-default-header-args:jupyter-julia '((:async . "yes")
;;                                                     (:session . "jupyter-julia")
;;                                                     (:kernel . "julia-1.6")
;; 						    (:exports . "both")
;; 						    (:eval . "never-export")))


;; https://datatofish.com/add-julia-to-jupyter/
;;  for setup
;; (add-to-list 'org-babel-load-languages '(julia . t) t)
;; (setq inferior-julia-program-name "/usr/local/bin/julia")

;; ** R setup

;; https://developers.refinitiv.com/en/article-catalog/article/setup-jupyter-notebook-r
;; (setq org-babel-default-header-args:jupyter-R
;;       '((:results . "value")
;; 	(:session . "jupyter-R")
;; 	(:kernel . "ir")
;; 	(:pandoc . "t")
;; 	(:exports . "both")
;; 	(:cache .   "no")
;; 	(:noweb . "no")
;; 	(:hlines . "no")
;; 	(:tangle . "no")
;; 	(:eval . "never-export")))

;; (require 'jupyter-R)

;; ** do the languages setup

(org-babel-do-load-languages 'org-babel-load-languages org-babel-load-languages)

;; * syntax highlighting
(add-to-list 'org-src-lang-modes '("jupyter-python" . python))
;; (add-to-list 'org-src-lang-modes '("jupyter-julia" . julia))
;; (add-to-list 'org-src-lang-modes '("jupyter-R" . R))

;; * make old ipython blocks work with jupyter python

(defalias 'org-babel-execute:ipython 'org-babel-execute:jupyter-python)
(setq org-babel-default-header-args:ipython org-babel-default-header-args:jupyter-python)
(add-to-list 'org-src-lang-modes '("ipython" . python))


;; * Kill the kernel when the buffer is closed
;; This is also harder than needed I think. This should kill the kernel in the current file.
(defun scimax-jupyter-get-session ()
  "Get the session name in the current buffer."
  (let ((lang (car (org-babel-get-src-block-info))))
    (or
     (cdr
      (assoc :session
	     (cadr (org-babel-params-from-properties lang))))
     (cdr (assoc :session
		 org-babel-default-header-args:jupyter-python)))))


(defun scimax-jupyter-org-kill-kernel ()
  "Kill the kernel."
  (interactive)
  (let ((session (scimax-jupyter-get-session)))
    (catch 'done
      (cl-loop for buffer in (buffer-list)
	       do
	       (when (string-match
		      (format "*jupyter-repl\\[.*\\]-%s*" session)
		      (buffer-name buffer))
		 (kill-buffer buffer)
		 (setq header-line-format nil)
		 (throw 'done buffer))))))


(defun scimax-jupyter-kill-kernel-hook (&rest args)
  (setq header-line-format
	(format "%s running. Click here to kill it."
		(scimax-jupyter-get-session)))
  (local-set-key [header-line down-mouse-1]
		 `(lambda ()
		    (interactive)
		    (scimax-jupyter-org-kill-kernel)))
  (add-hook 'kill-buffer-hook 'scimax-jupyter-org-kill-kernel nil t))


(defun scimax-jupyter-check-restart (&rest args)
  "If :restart is in the header, kill the kernel first."
  (let ((src (org-element-context)))
    (when (string-match ":restart" (or (org-element-property :parameters src) ""))
      (scimax-jupyter-org-kill-kernel))))


(defcustom scimax-jupyter-advices
  '((org-babel-execute:jupyter :before scimax-jupyter-kill-kernel-hook)
    (org-babel-execute:jupyter :before scimax-jupyter-check-restart)
    (jupyter-org-sync-results :override scimax-jupyter-org-sync-results)
    (jupyter-org--add-result :override scimax-jupyter-org--add-result)
    (jupyter-org-export-block-or-pandoc :override scimax-jupyter-org-export-block-or-pandoc))
  "Advices for scimax-jupyter.
This is a list of (emacs-jupyter-fn :position scimax-jupyter-fn)"
  :group 'scimax-jupyter)


(defun scimax-jupyter-advise ()
  "Turn scimax-jupyter advices on."
  (interactive)
  (cl-loop for (emacs-jupyter-fn position scimax-jupyter-fn) in scimax-jupyter-advices
	   do
	   (advice-add emacs-jupyter-fn position scimax-jupyter-fn)))


(defun scimax-jupyter-unadvise ()
  "Turn scimax-jupyter advices off."
  (interactive)
  (message "Un-advising emacs-jupyter")
  
  (cl-loop for (emacs-jupyter-fn position scimax-jupyter-fn) in scimax-jupyter-advices
	   do
	   (advice-remove emacs-jupyter-fn scimax-jupyter-fn)))


;; Turn on by default
(scimax-jupyter-advise)

;; * Fixing raw results
;; in emacs-jupyter :results raw, :results drawer did not work. This function seems to do the trick.

;; This function in emacs-jupyter seems to take a long time, and maybe is run
;; asynchronously. That messes up the order of results, as faster running things
;; come first in some process filter. This avoids this step and we do it later
;; in `scimax-jupyter-org-sync-results'. There is probably no guarantee that
;; this always works though.
(defun scimax-jupyter-org-export-block-or-pandoc (type value params)
  "Return VALUE, either converted with pandoc or in an export block.
If PARAMS has non-nil value for key ':pandoc' and TYPE is in
`jupyter-org-pandoc-convertable', convert the result with pandoc.
Otherwise, wrap it in an export block."
  (jupyter-org-export-block type value))


;; used to advise `jupyter-org-sync-results'. I do this to fix some issues like
;; results drawers, etc.
(defun scimax-jupyter-org-sync-results (req)
  "Return the result string in org syntax for the results of REQ.
Meant to be used as the return value of
`org-babel-execute:jupyter'."
  (when-let* ((results (nreverse (jupyter-org-request-results req))) 
	      ;; Not sure why these are reversed above, it works right for some things,
	      ;; but not right for display stuff. I wonder if there is some way
	      ;; to tell when something is displayed.
              (params (jupyter-org-request-block-params req))
              (result-params (alist-get :result-params params)))

    ;; There is an inconsistency when you specify :results value where printed
    ;; outputs and the return value is provided, sometimes as a fixed width.
    ;; That is consistent with the jupyter notebook, but not consistent with
    ;; org-babel. Here I only return the last thing if value is selected, and do
    ;; not require it to be fixed-width. That also isn't quite right, except in
    ;; a drawer. 
    (when (member "value" result-params)
      (let ((retval (car (last results)))) 
	(setq results
	      (list
	       (cond
		((eq 'fixed-width (org-element-type retval))
		 (org-element-property :value retval))
		(t
		 retval))))))

    (when (member "both" result-params)
      (let ((retval (car (last results)))) 
	(setf (car (last results))
	      (cond
	       ((eq 'fixed-width (org-element-type retval))
		(org-element-property :value retval))
	       (t
		retval)))))

    ;; Do pandoc conversion here so the order of results is preserved.
    (when (cdr (assoc :pandoc params))
      (setq results
	    (cl-loop for result in results
		     collect
		     (cond
		      ((stringp result)
		       result)
		      ((and (listp result)
			    (member (org-element-property :type result) jupyter-org-pandoc-convertable))
		       (jupyter-pandoc-convert
			(org-element-property :type result)
			"org"
			(org-element-property :value result)))
		      (t
		       result)))))
    
    (org-element-interpret-data
     (cond
      ;; This happens when a named block is a variable in another block.
      ;; It is different than a :results silent header.
      ((jupyter-org-request-silent-p req)
       results)

      ((member "raw" result-params)
       results)

      ;; fall through to a drawer for now.
      (t
       (apply #'jupyter-org-results-drawer results))))))



;; This was causing a problem in emacs-jupyter for using code blocks as
;; variables in other blocks. Commenting out the first line seems to fix it for
;; me.
(defun scimax-jupyter-org--add-result (req result)
  (cond
   ;; ((jupyter-org-request-silent-p req)
   ;;  (unless (equal (jupyter-org-request-silent-p req) "none")
   ;;    (message "%s" (org-element-interpret-data result))))
   ((jupyter-org-request-async-p req)
    (jupyter-org--clear-request-id req)
    (jupyter-org--do-insert-result req result))
   (t
    (push result (jupyter-org-request-results req)))))


;; I don't know where the \\ lines come from, this removes them.
(defun scimax-rm-backslashes ()
  "rm \\ from the end of lines."
  (let ((result-start (org-babel-where-is-src-block-result)))
    (save-excursion
      (when (and result-start
		 (goto-char result-start)
		 (looking-at org-babel-result-regexp))
	(while (re-search-forward "\\\\\\\\" (org-babel-result-end) t)
	  (replace-match "")
	  ;; I don't know why this is needed, but without, it was not finding everything.
	  (goto-char result-start))))))

(add-hook 'org-babel-after-execute-hook 'scimax-rm-backslashes)

;; * Fix org-show-entry so it doesn't collapse results all the time
;; 
;; I like to see the results, and this function usually hides them. I added the
;; unless wrapper, so it doesn't happen if point is on a src block. In a
;; previous version I just commented out the org-cycle-hide-drawers line, and I
;; may need to go back to that. this version should work as the original one
;; though except in src blocks. normally it is triggered by jumping to the next
;; src-block, and this is when I don't want it to hide the previous results.

(defun scimax-org-show-entry ()
  "Show the body directly following this heading.
Show the heading too, if it is currently invisible."
  (interactive)
  (unless (org-in-src-block-p)
    (save-excursion
      (ignore-errors
	(org-back-to-heading t)
	(outline-flag-region
	 (max (point-min) (1- (point)))
	 (save-excursion
	   (if (re-search-forward
		(concat "[\r\n]\\(" org-outline-regexp "\\)") nil t)
	       (match-beginning 1)
	     (point-max)))
	 nil)
	(org-cycle-hide-drawers 'children)))))

(advice-add 'org-show-entry :override #'scimax-org-show-entry)

;; * Working with exceptions
;;
;; The syntax highlighting for exceptions seems broken. `jupyter-org-client' has
;; a bunch of code that suggests this should be working, even referencing my
;; blog post about it. This is ok for now I guess. As I recall the syntax
;; highlighting is not persistent if you don't put the ansi codes in the buffer.
;;
;; this handled in generic functions that I find too difficult to trace for now.

(defun scimax-jupyter-jump-to-error ()
  "In a src block, jump to the line indicated as an error in the results.
In a SyntaxError, there is not a traceback with a line number, so
we handle it separately. It doesn't seem like it should be that
way, but it is."
  (interactive)
  (let* ((cp (point))
	 (location (org-babel-where-is-src-block-result))
	 (case-fold-search t))
    
    (when (and location
	       (goto-char location)
	       (looking-at org-babel-result-regexp))
      (cond
       ;; Check for SyntaxError
       ((string-match "SyntaxError:" (buffer-substring location (org-babel-result-end)))
	(re-search-forward (rx (zero-or-more " ") "^") nil (org-babel-result-end))
	(previous-line)
	(let ((pattern (string-trim-left
			(buffer-substring-no-properties
			 (line-beginning-position) (line-end-position)))))
	  (goto-char cp)
	  (goto-char (org-element-property :begin (org-element-context)))
	  (unless
	      (search-forward pattern (org-element-property :end (org-element-context)) t)
	    (message "No SyntaxError found like %s" pattern))))

       ;; search for something like --> 21
       (t
	(goto-char location)
	(re-search-forward "^-*> \\([[:digit:]]*\\)" (org-babel-result-end))
	(save-match-data
	  (goto-char cp)
	  (goto-char (org-element-property :begin (org-element-context))))
	(forward-line (string-to-number (match-string-no-properties 1))))))))


;; * Handling ansi codes

(defun scimax-jupyter-ansi ()
  "Replaces ansi-codes in exceptions with colored text.
I thought emacs-jupyter did this automatically, but it may only
happen in the REPL. Without this, the tracebacks are very long
and basically unreadable.

We also add some font properties to click on goto-error.

This should only apply to jupyter-lang blocks."
  (when (string-match "^jupyter" (car (or (org-babel-get-src-block-info t) '(""))))
    (let* ((r (org-babel-where-is-src-block-result))
	   (result (when r
		     (save-excursion
		       (goto-char r)
		       (org-element-context)))))
      (when result
	(ansi-color-apply-on-region (org-element-property :begin result)
				    (org-element-property :end result))

	;; Let's fontify "# [goto error]" to it is clickable
	(save-excursion
	  (goto-char r)
	  (when (search-forward "# [goto error]" (org-element-property :end result) t)
	    (add-text-properties
	     (match-beginning 0) (match-end 0)
	     (list 'help-echo "Click to jump to error."
		   'mouse-face 'highlight
		   'local-map (let ((map (copy-keymap help-mode-map)))
				(define-key map [mouse-1] (lambda ()
							    (interactive)
							    (search-backward "#+BEGIN_SRC")
							    (scimax-jupyter-jump-to-error)))
				map))))))
      
      t)))


(add-to-list 'org-babel-after-execute-hook 'scimax-jupyter-ansi t)

;; * The scimax jupyter hydra
;; customization of what is in jupyter
;; These are more aligned with jupyter notebook I think

;; I had to use pretty hydra to get nicely aligned columns here.
(use-package pretty-hydra)

(pretty-hydra-define scimax-jupyter-org-hydra (:color blue :hint nil)
  ("Execute"
   (("<return>" org-ctrl-c-ctrl-c "current" :color red)
    ("C-<return>" org-ctrl-c-ctrl-c "current" :color red)
    ("S-<return>" jupyter-org-execute-and-next-block "current and next" :color red)
    ("M-<return>" (progn (org-ctrl-c-ctrl-c) (scimax-ob-insert-src-block t)) "current and new" )
    ("C-M-<return>" jupyter-org-execute-subtree "subtree")
    ("S-C-<return>" jupyter-org-restart-and-execute-to-point "restart to point")
    ("S-M-<return>" jupyter-org-restart-kernel-execute-buffer "restart buffer"))

   "Navigate"
   (("p" org-babel-previous-src-block "previous" :color red)
    ("P" jupyter-org-previous-busy-src-block "previous busy")
    ("n" org-babel-next-src-block  "next" :color red)
    ("N" jupyter-org-next-busy-src-block "next busy" :color red)
    ("g" jupyter-org-jump-to-visible-block "jump to visible src")
    ("G" jupyter-org-jump-to-block "jump to src block")
    ("e" scimax-jupyter-jump-to-error "Jump to error"))

   "Edit"
   (("<up>" jupyter-org-move-src-block "move up" :color red)
    ("<down>" (jupyter-org-move-src-block t) "move down" :color red)
    ("x" jupyter-org-kill-block-and-results "kill block")
    ("c" jupyter-org-copy-block-and-results "copy block")
    ("o" (jupyter-org-clone-block t) "clone")
    ("m" jupyter-org-merge-blocks "merge")
    ("s" jupyter-org-split-src-block "split")
    ("a" (jupyter-org-insert-src-block nil current-prefix-arg) "insert above")
    ("b" (jupyter-org-insert-src-block t current-prefix-arg) "insert below")
    ("l" org-babel-remove-result "clear result")
    ("L" jupyter-org-clear-all-results "clear all results")
    ("h" jupyter-org-edit-header "edit header"))

   "Misc"
   (("i" jupyter-org-inspect-src-block "inspect")
    ("<tab>" completion-at-point "Complete")
    
    ("O" scimax-ob/body "scimax-ob")
    ("q" nil "quit"))
   
   "Kernel"
   (("s" org-babel-jupyter-scratch-buffer "scratch")
    ("z" org-babel-switch-to-session "REPL")
    ("u" jupyter-org-interrupt-kernel "interrupt")
    ("r" (jupyter-org-with-src-block-client
	  (jupyter-repl-restart-kernel)) "restart")
    ("k" scimax-jupyter-org-kill-kernel "kill"))))


(jupyter-org-define-key (kbd "<f12>") #'scimax-jupyter-org-hydra/body)

;; utilities
;; * autopep8

(defun autopep8 ()
  "Replace Python code block contents with autopep8 corrected code."
  (interactive)
  (unless (executable-find "autopep8")
    (if (executable-find "pip")
	(shell-command "python -c \"import pip; pip.main(['install','autopep8'])\"")
      (shell-command "python -c \"from setuptools.command import easy_install; easy_install.main(['-U','autopep8'])\"")))
  (let* ((src (org-element-context))
	 (beg (org-element-property :begin src))
	 (value (org-element-property :value src)))
    (save-excursion
      (goto-char beg)
      (search-forward value)
      (shell-command-on-region
       (match-beginning 0)
       (match-end 0)
       "autopep8 -a -a -" nil t))))


;; * pylint
(defvar pylint-options
  '()
  "List of options to use with pylint.")

(setq pylint-options
      '("-r no "			 ; no reports
	;; we are not usually writing programs where it
	;; makes sense to be too formal on variable
	;; names.
	"--disable=invalid-name "
	;; don't usually have modules, which triggers
	;; this when there is not string at the top
	"--disable=missing-docstring "
	;; superfluous-parens is raised with print(),
	;; which I am promoting for python3
	;; compatibility.
	"--disable=superfluous-parens "	;

	;; these do not seem important for my work.
	"--disable=too-many-locals "	;

	;; this is raised in solving odes and is
	;; unimportant for us.
	"--disable=unused-argument "	;
	"--disable=unused-wildcard-import "
	"--disable=redefined-outer-name "
	;; this is triggered a lot from fsolve
	"--disable=unbalanced-tuple-unpacking "
	"--disable=wildcard-import "
	"--disable=redefined-builtin "
	;; I dont mind semicolon separated lines
	"--disable=multiple-statements "
	;; pylint picks up np.linspace as a no-member error. That does not make sense.
	"--disable=no-member "
	"--disable=wrong-import-order "
	"--disable=unused-import "))


(defun pylint ()
  "Run pylint on a source block.
Opens a buffer with links to what is found. This function installs pylint if needed."
  (interactive)
  (let ((eop (org-element-at-point))
	(temporary-file-directory ".")
        (cb (current-buffer))
	(n) ; for line number
	(cn) ; column number
	(content) ; error on line
	(pb "*pylint*")
	(link)
	(tempfile))

    (unless (executable-find "pylint")
      (if (executable-find "pip")
	  (shell-command "python -c \"import pip; pip.main(['install','pylint'])\"")
	(shell-command "python -c \"from setuptools.command import easy_install; easy_install.main(['pylint'])\"")))

    ;; rm buffer if it exists
    (when (get-buffer pb) (kill-buffer pb))

    ;; only run if in a python code-block
    (when (and (eq 'src-block (car eop))
	       (string= "python" (org-element-property :language eop)))

      ;; tempfile for the code
      (setq tempfile (make-temp-file "org-py-check" nil ".py"))
      ;; create code file
      (with-temp-file tempfile
	(insert (org-element-property :value eop)))

      ;; pylint
      (let ((status (shell-command
		     (concat
		      "pylint "
		      (mapconcat 'identity pylint-options " ")
		      " "
		      ;; this is the file to check.
		      (file-name-nondirectory tempfile))))

	    ;; remove empty strings
	    (output (delete "" (split-string
				(with-current-buffer "*Shell Command Output*"
				  (buffer-string)) "\n"))))

	;; also remove this line so the output is empty if nothing
	;; comes up
	(setq output (delete
		      "No config file found, using default configuration"
		      output))

	(kill-buffer "*Shell Command Output*")
	(if output
	    (progn
	      (set-buffer (get-buffer-create pb))
	      (insert (format "\n\n* pylint (status = %s)\n" status))
	      (insert "pylint checks your code for errors, style and convention. Click on the links to jump to each line.

")

	      (dolist (line output)
		;; pylint gives a line and column number
		(if
		    (string-match "[A-Z]:\\s-+\\([0-9]*\\),\\s-*\\([0-9]*\\):\\(.*\\)"
				  line)
		    (let ((line-number (match-string 1 line))
			  (column-number (match-string 2 line))
			  (content (match-string 3 line)))

		      (setq link (format "[[elisp:(progn (switch-to-buffer-other-window \"%s\")(goto-char %s)(forward-line %s)(forward-line 0)(forward-char %s))][%s]]\n"
					 cb
					 (org-element-property :begin eop)
					 line-number
					 column-number
					 line)))
		  ;; no match, just insert line
		  (setq link (concat line "\n")))
		(insert link)))
	  (message "pylint was clean!")))

      (when (get-buffer pb)
	;; open the buffer
	(switch-to-buffer-other-window pb)
	(goto-char (point-min))
	(insert "Press q to close the window\n")
	(org-mode)
	(org-cycle '(64))  ; open everything
	;; make read-only and press q to quit
	(setq buffer-read-only t)
	(use-local-map (copy-keymap org-mode-map))
	(local-set-key "q" #'(lambda () (interactive) (kill-buffer)))
	(switch-to-buffer-other-window cb))
      ;; final cleanup and delete file
      (delete-file tempfile))))


;; * Numbered lines in code blocks
(defvar number-line-overlays '()
  "List of overlays for line numbers.")

(make-variable-buffer-local 'number-line-overlays)

(defun number-line-src-block ()
  "Add line numbers to an org src-block."
  (interactive)
  (save-excursion
    (let* ((src-block (org-element-context))
           (nlines (- (length
                       (s-split
                        "\n"
                        (org-element-property :value src-block)))
                      1)))
      ;; clear any existing overlays
      (when number-line-overlays
	(mapc 'delete-overlay
	      number-line-overlays)
	(setq number-line-overlays '()))

      (goto-char (org-element-property :begin src-block))
      ;; the beginning may be header, so we move forward to get the #+BEGIN
      ;; line. Then jump one more to get in the code block
      (while (not (looking-at "#\\+BEGIN"))
	(forward-line))
      (forward-line)
      (cl-loop for i from 1 to nlines
               do
               (beginning-of-line)
               (let (ov)
		 (setq ov (make-overlay (point)(point)))
		 (overlay-put
		  ov
		  'before-string (propertize
				  (format "%03s:" (number-to-string i))
				  'font-lock-face '(:foreground "black" :background "gray80")
				  'local-map (let ((map (make-sparse-keymap)))
					       (define-key map [mouse-1]
						 (lambda ()
						   (interactive)
						   (mapc 'delete-overlay
							 number-line-overlays)
						   (setq number-line-overlays '())))
					       map)))
		 (overlay-put ov 'mouse-face 'highlight)
		 (overlay-put ov 'help-echo "Click to remove")
		 (overlay-put ov 'local-map (let ((map (make-sparse-keymap)))
					      (define-key map [mouse-1]
						(lambda ()
						  (interactive)
						  (mapc 'delete-overlay
							number-line-overlays)
						  (setq number-line-overlays '())))
					      map))
		 (add-to-list 'number-line-overlays ov))
               (forward-line))))
  (add-hook 'post-command-hook 'number-line-src-block nil 'local))

(provide 'scimax-jupyter)

;;; scimax-jupyter.el ends here