Add jupyter-org-display-execution-time functionality

* jupyter-base.el (jupyter-format-time)
(jupyter-execution-time): New functions.

* jupyter-org-client.el (jupyter-org-display-execution-time): New
custom variable.
(jupyter-generate-request): Clear any previous execution time strings
before generating the request for the source block.
(jupyter-org--display-execution-time)
(jupyter-org-clear-execution-time): New functions.
(jupyter-handle-execute-reply): Display execution time on the source
block associated with a request.

Author: nnicandro

jupyter-base.el

@@ -497,6 +497,32 @@ call the handler methods of those types."
     ("interrupt_request" "control")
     (_ "shell")))
 
+(declare-function jupyter-message-time "jupyter-messages")
+(declare-function jupyter-find-message "jupyter-monads")
+
+(defun jupyter-execution-time (req)
+  "Return an approximate value of the execution time of REQ.
+Return a time value in the same form as what is returned by
+`current-time' or nil when a value could not be computed,
+e.g. when REQ does not represent an execute_request or has not
+completed yet.
+
+Note this just returns the time difference between the
+execute_input message of the request and the execute_reply
+message.  The execute_input message is received when the kernel
+begins executing the contents of REQ and the execute_reply is
+received when the kernel has completed execution of REQ minus how
+long it actually takes for the request to generate all of its
+potential output.  So the returned time, in addition to the true
+execution time of the code, contains the time taken in preparing
+the request to be executed and other kernel specific tasks."
+  (when-let* ((msgs (jupyter-request-messages req))
+              (ex-input (jupyter-find-message "execute_input" msgs))
+              (ex-reply (jupyter-find-message "execute_reply" msgs)))
+    (time-subtract
+     (jupyter-message-time ex-reply)
+     (jupyter-message-time ex-input))))
+
 ;;; Connecting to a kernel's channels
 
 (eval-when-compile (require 'tramp))
@@ -768,6 +794,48 @@ TIME is assumed to have the same form as the return value of
                   minutes
                   (if past " ago" ""))))))))
 
+(defun jupyter-format-time (time)
+  "Return a description string describing TIME.
+TIME is a time value as returned by `current-time'.  Return
+strings like
+
+    \"1 day 5 min 10.330 s\", \"10 min 5.000 s\", \"30.200 s\"
+
+depending on the number of seconds contained in TIME.  If TIME is
+nil, return \"Never\"."
+  (if (null time) "Never"
+    (let* ((result "")
+           (s (float-time time))
+           (d (floor (/ s 86400.0))))
+      (unless (zerop d)
+        (cl-callf concat result
+          (format "%d day%s" d (if (= d 1) "" "s"))))
+      (let ((h (floor (- (/ s 3600.0) (* 24 d)))))
+        (unless (zerop h)
+          (cl-callf concat result
+            (format "%s%d hour%s"
+                    (if (zerop (length result)) "" " ")
+                    h
+                    (if (= h 1) "" "s"))))
+        (let ((m (floor (- (/ s 60.0)
+                           (+ (* 60 h)
+                              (* 1440 d))))))
+          (unless (zerop m)
+            (cl-callf concat result
+              (format "%s%d min%s"
+                      (if (zerop (length result)) "" " ")
+                      m
+                      (if (= m 1) "" "s"))))
+          (setq s (- s (+ (* 60 m)
+                          (* 3600 h)
+                          (* 86400 d))))
+          (unless (zerop s)
+            (cl-callf concat result
+              (format "%s%.3f s"
+                      (if (zerop (length result)) "" " ")
+                      s)))))
+      result)))
+
 ;;; Simple weak references
 ;; Thanks to Chris Wellon https://nullprogram.com/blog/2014/01/27/
 

jupyter-org-client.el

@@ -91,6 +91,20 @@ after the failed one are not executed."
   :group 'ob-jupyter
   :type 'boolean)
 
+(defcustom jupyter-org-display-execution-time nil
+  "Whether or not to display the execution time of a source block.
+If this variable is nil, the execution times are not displayed.
+When it is t, display the execution times regardless of how long
+it took to execute.  When it is a number, display the execution
+time when it is longer than that many seconds.
+
+To clear the execution time information from the source block
+simply edit it or call `jupyter-org-clear-execution-time'."
+  :group 'ob-jupyter
+  :type '(choice (const :tag "Never display" nil)
+                 (const :tag "Always display" t)
+                 (number :tag "Display when above this threshold (in seconds)")))
+
 (defcustom jupyter-org-resource-directory "./.ob-jupyter/"
   "Directory used to store automatically generated image files.
 See `jupyter-org-image-file-name'."
@@ -216,6 +230,7 @@ nothing and return nil."
       ;; started due to sending a completion request.
       (save-excursion
         (goto-char org-babel-current-src-block-location)
+        (jupyter-org-clear-execution-time)
         (let* ((context (org-element-context))
                (block-params org-babel-jupyter-current-src-block-params)
                (result-params (alist-get :result-params block-params))
@@ -422,7 +437,51 @@ to."
       (forward-line)
       (insert (org-element-normalize-string (plist-get pl :text))))))
 
+(defun jupyter-org--display-execution-time (req)
+  "In the Org buffer of REQ, show the REQ's execution time."
+  (pcase jupyter-org-display-execution-time
+    ((and (or (and `t
+                   (let time (jupyter-execution-time req)))
+              (and (pred numberp) secs
+                   (let time (jupyter-execution-time req))
+                   (guard (> time secs))))
+          (let (cl-struct jupyter-org-request inline-block-p) req)
+          (guard (not inline-block-p)))
+     (jupyter-org-with-point-at req
+       (let* ((src-block (org-element-at-point))
+              (ov (make-overlay
+                   (org-element-property :begin src-block)
+                   ;; Exclude the newline to make it look like
+                   ;;
+                   ;;    #+end_src Execution time ...
+                   (1- (jupyter-org-element-end-before-blanks src-block)))))
+         (let ((delete
+                (list (lambda (&rest _)
+                        (delete-overlay ov)))))
+           (overlay-put ov 'evaporate t)
+           (overlay-put ov 'jupyter-execution-time time)
+           (overlay-put ov 'modification-hooks delete)
+           (overlay-put ov 'insert-in-front-hooks delete)
+           (overlay-put ov 'insert-behind-hooks delete))
+         (overlay-put
+          ov 'after-string
+          (concat " " (propertize
+                       (format "Execution time: %s"
+                               (jupyter-format-time time))
+                       'face 'bold-italic))))))))
+
+(defun jupyter-org-clear-execution-time ()
+  "Clear the execution time overlay for the source block at point."
+  (interactive)
+  (let ((el (org-element-at-point)))
+    (pcase (org-element-type el)
+      ((or `src-block `babel-call)
+       (dolist (ov (overlays-at (org-element-property :begin el)))
+         (when (overlay-get ov 'jupyter-execution-time)
+           (delete-overlay ov)))))))
+
 (cl-defmethod jupyter-handle-execute-reply ((_client jupyter-org-client) (req jupyter-org-request) msg)
+  (jupyter-org--display-execution-time req)
   (jupyter-with-message-content msg (status payload)
     (when payload
       (jupyter-org-with-point-at req