Implement kernel-backed completion

Author: Greg Sexton

README.org

@@ -157,19 +157,24 @@
      SRC block, place the point on the thing you're interested in and
      run =M-x ob-ipython-inspect=. I recommend you bind this to a key.
 
+   * ~ob-ipython-completions~ queries the kernel for completions at a
+     position. You may use this to hook up any completion mechanism.
+     We already provide a company backend. With company installed, add
+     ~(add-to-list 'company-backends 'company-ob-ipython)~ somewhere
+     in your config. This should then work while editing a src block.
+
    * It's often easier to play with code using a REPL. With the point
      in an ipython SRC block, you can open a REPL connected to the
-     current kernel by running =C-c C-v C-z=. I recommend you do this
-     anyway, as python-mode can now use this REPL to provide
-     completion in code buffers.
+     current kernel by running =C-c C-v C-z=.
 
    * If evaluated code produces an error, this will be displayed
      nicely in a buffer using IPython's traceback support.
 
-   * Stdout from code evaluation is displayed in a popup buffer. This
-     is great for debugging or getting verbose output that is best
-     left out of documents. If you wish to capture stdout in your
-     document use the =:results output= SRC block header.
+   * Stdout/err from code evaluation is displayed in a popup buffer.
+     This is great for debugging or getting verbose output that is
+     best left out of documents (e.g progress updates). If you wish to
+     capture output in your document use the =:results output= SRC
+     block header.
 
    * You can interrupt or kill a running kernel. This is helpful if
      things get stuck or really broken. See =M-x
@@ -214,8 +219,6 @@
        (add-hook 'org-babel-after-execute-hook 'org-display-inline-images 'append)
      #+END_SRC
 
-   * Open a REPL using =C-c C-v C-z= so that you get completion in Python buffers.
-
    * Export with the =LaTeX= backend using the =minted= package for
      source block highlighting fails for =ipython= blocks by default
      with the error

client.py

@@ -9,26 +9,31 @@
 semaphore = threading.Semaphore(value=0)
 interested_lock = threading.Lock()
 interested = []
-hasreply = False
-hasidle = False
 
-def msg_router(name, ch):
-    global hasreply, hasidle
+def msg_router(io, shell):
     while True:
-        msg = ch()
-        msg['channel'] = name
+        msg = io()
+        msg['channel'] = 'io'
         msgid = msg['parent_header'].get('msg_id', None)
         with interested_lock:
             if msgid not in interested:
                 continue
-        if msg.get('msg_type', '') in ['execute_reply', 'inspect_reply']:
-            hasreply = True
-        elif (msg.get('msg_type', '') == 'status' and
+        print(json.dumps(msg, default=str))
+        if (msg.get('msg_type', '') == 'status' and
             msg['content']['execution_state'] == 'idle'):
-            hasidle = True
-        if not msg['msg_type'] in ['status', 'execute_input']:
-            print(json.dumps(msg, default=str))
-        if hasreply and hasidle:
+            break
+
+    while True:
+        msg = shell()
+        msg['channel'] = 'shell'
+        msgid = msg['parent_header'].get('msg_id', None)
+        with interested_lock:
+            if msgid not in interested:
+                continue
+        print(json.dumps(msg, default=str))
+        if msg.get('msg_type', '') in ['execute_reply',
+                                       'inspect_reply',
+                                       'complete_reply']:
             semaphore.release()
 
 def create_client(name):
@@ -39,17 +44,17 @@ def create_client(name):
     c = client.BlockingKernelClient(connection_file=cf)
     c.load_connection_file()
     c.start_channels()
-    chans = [('io', c.get_iopub_msg), ('shell', c.get_shell_msg), ('stdin', c.get_stdin_msg)]
-    for name, ch in chans:
-        t = threading.Thread(target=msg_router, args=(name, ch))
-        t.setDaemon(True)
-        t.start()
+    io, shell = c.get_iopub_msg, c.get_shell_msg
+    t = threading.Thread(target=msg_router, args=(io, shell))
+    t.setDaemon(True)
+    t.start()
     return c
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--conn-file')
 parser.add_argument('--execute', action='store_true')
 parser.add_argument('--inspect', action='store_true')
+parser.add_argument('--complete', action='store_true')
 args = parser.parse_args()
 
 c = create_client(args.conn_file)
@@ -67,4 +72,16 @@ def create_client(name):
                           detail_level=req.get('detail', 0))
         interested.append(msgid)
 
+    elif args.complete:
+        req = json.loads(sys.stdin.read())
+        code = req['code']
+        pos = req.get('pos', len(code))
+        # causes things to hang as kernel doesn't come back with a
+        # complete_reply
+        if code[pos-1] in ['\n', '\r']:
+            sys.exit(0)
+        msgid = c.complete(code,
+                           cursor_pos=pos)
+        interested.append(msgid)
+
 semaphore.acquire()

ob-ipython.el

@@ -317,7 +317,7 @@ a new kernel will be started."
 
 (defun ob-ipython--extract-status (msgs)
   (->> msgs
-       (-filter (lambda (msg) (-contains? '("execute_reply" "inspect_reply")
+       (-filter (lambda (msg) (-contains? '("execute_reply" "inspect_reply" "complete_reply")
                                           (cdr (assoc 'msg_type msg)))))
        car
        (assoc 'content)
@@ -340,7 +340,7 @@ a new kernel will be started."
     (with-temp-buffer
       (let ((ret (apply 'call-process-region input nil
                         (ob-ipython--get-python) nil t nil
-                        ;; TODO: hardcoded default
+                        ;; TODO: hardcoded default -- can use local org-src--babel-info
                         (list "--" ob-ipython-client-path "--conn-file" "default" "--inspect"))))
         (if (> ret 0)
             (ob-ipython--dump-error (buffer-string))
@@ -368,19 +368,76 @@ a new kernel will be started."
   "Ask a kernel for documentation on the thing at POS in BUFFER."
   (interactive (list (current-buffer) (point)))
   (-if-let (result (->> (ob-ipython--inspect buffer pos) (assoc 'text/plain) cdr))
-      (ob-ipython--create-inspect-buffer result)
-    (message "No documentation was found.")))
+    (ob-ipython--create-inspect-buffer result)
+  (message "No documentation was found.")))
+
+;; completion
+
+(defun ob-ipython--complete-request (code &optional pos)
+  (let ((input (json-encode `((code . ,code)
+                              (pos . ,(or pos (length code)))))))
+    (with-temp-buffer
+      (let ((ret (apply 'call-process-region input nil
+                        (ob-ipython--get-python) nil t nil
+                        ;; TODO: hardcoded default
+                        (list "--" ob-ipython-client-path "--conn-file" "default" "--complete"))))
+        (if (> ret 0)
+            (ob-ipython--dump-error (buffer-string))
+          (goto-char (point-min))
+          (ob-ipython--collect-json))))))
+
+(defun ob-ipython-completions (buffer pos)
+  "Ask a kernel for completions on the thing at POS in BUFFER."
+  (let* ((code (with-current-buffer buffer
+                 (buffer-substring-no-properties (point-min) (point-max))))
+         (resp (ob-ipython--complete-request code pos))
+         (status (ob-ipython--extract-status resp)))
+    (if (not (string= "ok" status))
+        '()
+      (->> resp
+           (-filter (lambda (msg)
+                      (-contains? '("complete_reply")
+                                  (cdr (assoc 'msg_type msg)))))
+           (-mapcat (lambda (msg)
+                      (->> msg
+                           (assoc 'content)
+                           cdr)))))))
+
+(defun company-ob-ipython (command &optional arg &rest ignored)
+  (interactive (list 'interactive))
+  (cl-case command
+    (interactive (company-begin-backend 'company-ob-ipython))
+    (prefix (and
+             ob-ipython-mode
+             (let ((res (ob-ipython-completions (current-buffer) (1- (point)))))
+               (substring (buffer-string) (cdr (assoc 'cursor_start res))
+                          (cdr (assoc 'cursor_end res))))))
+    (candidates (let ((res (ob-ipython-completions (current-buffer) (1- (point)))))
+                  (cdr (assoc 'matches res))))
+    (sorted t)))
+
+;; mode
+
+(define-minor-mode ob-ipython-mode
+  ""
+  nil
+  " ipy"
+  '())
 
 ;; babel framework
 
 (add-to-list 'org-src-lang-modes '("ipython" . python))
 
 (defvar org-babel-default-header-args:ipython '())
 
+(defun org-babel-edit-prep:ipython (info)
+  ;; TODO: based on kernel, should change the mode
+  (ob-ipython-mode +1))
+
 (defun ob-ipython--normalize-session (session)
   (if (string= "default" session)
-      (error "default is reserved for when no name is provided. Please use a different session name.")
-    (or session "default")))
+    (error "default is reserved for when no name is provided. Please use a different session name.")
+  (or session "default")))
 
 (defun org-babel-execute:ipython (body params)
   "Execute a block of IPython code with Babel.