[PATCH] Expanded ob-python results handling and plotting

[Jack Kamm <jackkamm@gmail.com>]

[-- Attachment #1: Type: text/plain, Size: 1639 bytes --]

The attached patch adds ob-python value results handling for the
following types of results:

- Dictionaries
- Numpy arrays
- Pandas dataframes
- Matplotlib figures

This is a bigger commit than I'm used to, so I thought I better send it
out before merging, in case someone notices obvious problems I missed.

Overview of changes:

Dictionaries are now transformed into alists before being converted to
lisp. Previously, they had been getting mangled, like so:

#+begin_src python
  return {"a": 1, "b": 2}
#+end_src

#+RESULTS:
| a | : | 1 | b | : | 2 |

But now they appear like so:

#+begin_src python
  return {"a": 1, "b": 2}
#+end_src

#+RESULTS:
| a | 1 |
| b | 2 |

Numpy arrays and pandas dataframes are also converted to tables
automatically now. Tables converted from Pandas dataframes have row and
column names.

To avoid conversion, you can specify "raw", "verbatim", "scalar", or
"output" in the ":results" header argument.

For plotting, you can specify "graphics" in the ":results"
header. You'll also need to provide a ":file" argument. The behavior
depends on whether using output or value results. For output results,
the current figure (pyplot.gcf) is cleared before evaluating, then the
result saved. For value results, the block is expected to return a
matplotlib Figure, which is saved. To set the figure size, do it from
within Python.

Here is an example of how to plot:

#+begin_src python :results output graphics file :file boxplot.svg
  import matplotlib.pyplot as plt
  import seaborn as sns
  plt.figure(figsize=(5, 5))
  tips = sns.load_dataset("tips")
  sns.boxplot(x="day", y="tip", data=tips)
#+end_src


[-- Warning: decoded text below may be mangled, UTF-8 assumed --]
[-- Attachment #2: 0001-ob-python-Add-results-handling-for-dicts-dataframes-.patch --]
[-- Type: text/x-patch, Size: 11849 bytes --]

From 09f9c42bb629a356e1c36f04f69c8baf795b411b Mon Sep 17 00:00:00 2001
From: Jack Kamm <jackkamm@gmail.com>
Date: Tue, 25 Aug 2020 21:57:24 -0700
Subject: [PATCH] ob-python: Add results handling for dicts, dataframes,
 arrays, plots

* lisp/ob-python.el (org-babel-execute:python): Parse graphics-file
from params.
(org-babel-python--def-format-value): Python code for formatting
value results before returning.
(org-babel-python--output-graphics-wrapper): Python code for handling
output graphics results.
(org-babel-python--nonsession-value-wrapper): Replaces
org-babel-python-wrapper-method, org-babel-python-pp-wrapper-method.
(org-babel-python--session-output-wrapper): Renamed from
org-babel-python--exec-tmpfile.
(org-babel-python--session-value-wrapper): Renamed and modified from
org-babel-python--eval-ast.
(org-babel-python-evaluate-external-process): New parameter for
graphics file.
(org-babel-python-evaluate-session): New parameter for graphics file.

Added results handling for dictionaries, Pandas and numpy tables, and
matplotlib plots.
---
 etc/ORG-NEWS      |  16 +++++-
 lisp/ob-python.el | 122 +++++++++++++++++++++++++++++++---------------
 2 files changed, 96 insertions(+), 42 deletions(-)

diff --git a/etc/ORG-NEWS b/etc/ORG-NEWS
index 10658a970..4f9863a5b 100644
--- a/etc/ORG-NEWS
+++ b/etc/ORG-NEWS
@@ -66,8 +66,8 @@ to switch to the new signature.
 *** Python session return values must be top-level expression statements
 
 Python blocks with ~:session :results value~ header arguments now only
-return a value if the last line is a top-level expression statement.
-Also, when a None value is returned, "None" will be printed under
+return a value if the last line is a top-level expression statement,
+otherwise the result is None. Also, None will now show up under
 "#+RESULTS:", as it already did with ~:results value~ for non-session
 blocks.
 
@@ -235,6 +235,18 @@ Screen blocks now recognize the =:screenrc= header argument and pass
 its value to the screen command via the "-c" option. The default
 remains =/dev/null= (i.e. a clean screen session)
 
+*** =ob-python.el=: Support for more result types and plotting
+
+=ob-python= now recognizes dictionaries, numpy arrays, and pandas
+dataframes, and will convert them to org-mode tables when appropriate.
+
+When the header argument =:results graphic= is set, =ob-python= will
+use matplotlib to save graphics. The behavior depends on whether value
+or output results are used. For value results, the last line should
+return a matplotlib Figure object to plot. For output results, the
+current figure (as returned by =pyplot.gcf()=) is cleared and then
+plotted.
+
 *** =RET= and =C-j= now obey ~electric-indent-mode~
 
 Since Emacs 24.4, ~electric-indent-mode~ is enabled by default.  In
diff --git a/lisp/ob-python.el b/lisp/ob-python.el
index 44e1b63e0..92ca82625 100644
--- a/lisp/ob-python.el
+++ b/lisp/ob-python.el
@@ -79,6 +79,8 @@ (defun org-babel-execute:python (body params)
 	      org-babel-python-command))
 	 (session (org-babel-python-initiate-session
 		   (cdr (assq :session params))))
+	 (graphics-file (and (member "graphics" (assq :result-params params))
+			     (org-babel-graphical-output-file params)))
          (result-params (cdr (assq :result-params params)))
          (result-type (cdr (assq :result-type params)))
 	 (return-val (when (and (eq result-type 'value) (not session))
@@ -89,7 +91,8 @@ (defun org-babel-execute:python (body params)
 	   (concat body (if return-val (format "\nreturn %s" return-val) ""))
 	   params (org-babel-variable-assignments:python params)))
          (result (org-babel-python-evaluate
-		  session full-body result-type result-params preamble)))
+		  session full-body result-type result-params preamble
+		  graphics-file)))
     (org-babel-reassemble-table
      result
      (org-babel-pick-name (cdr (assq :colname-names params))
@@ -225,67 +228,102 @@ (defun org-babel-python-initiate-session (&optional session _params)
     (org-babel-python-session-buffer
      (org-babel-python-initiate-session-by-key session))))
 
-(defconst org-babel-python-wrapper-method
-  "
-def main():
+(defconst org-babel-python--def-format-value "\
+def __org_babel_python_format_value(result, result_file, result_params):
+    with open(result_file, 'w') as f:
+        if 'graphics' in result_params:
+            result.savefig(result_file)
+        elif 'pp' in result_params:
+            import pprint
+            f.write(pprint.pformat(result))
+        else:
+            if not set(result_params).intersection(\
+['scalar', 'verbatim', 'raw']):
+                def dict2alist(res):
+                    if isinstance(res, dict):
+                        return [(k, dict2alist(v)) for k, v in res.items()]
+                    elif isinstance(res, list) or isinstance(res, tuple):
+                        return [dict2alist(x) for x in res]
+                    else:
+                        return res
+                result = dict2alist(result)
+                try:
+                    import pandas as pd
+                except ModuleNotFoundError:
+                    pass
+                else:
+                    if isinstance(result, pd.DataFrame):
+                        result = [[''] + list(result.columns), None] + \
+[[i] + list(row) for i, row in result.iterrows()]
+                try:
+                    import numpy as np
+                except ModuleNotFoundError:
+                    pass
+                else:
+                    if isinstance(result, np.ndarray):
+                        result = result.tolist()
+            f.write(str(result))")
+
+(defun org-babel-python--output-graphics-wrapper
+    (body graphics-file)
+  "Wrap BODY to plot to GRAPHICS-FILE if it is non-nil."
+  (if graphics-file
+      (format "\
+import matplotlib.pyplot as __org_babel_python_plt
+__org_babel_python_plt.gcf().clear()
 %s
+__org_babel_python_plt.savefig('%s')" body graphics-file)
+    body))
 
-open('%s', 'w').write( str(main()) )")
-(defconst org-babel-python-pp-wrapper-method
-  "
-import pprint
+(defconst org-babel-python--nonsession-value-wrapper
+  (concat org-babel-python--def-format-value "
 def main():
 %s
 
-open('%s', 'w').write( pprint.pformat(main()) )")
+__org_babel_python_format_value(main(), '%s', %s)")
+  "TODO")
 
-(defconst org-babel-python--exec-tmpfile "\
+(defconst org-babel-python--session-output-wrapper "\
 with open('%s') as f:
     exec(compile(f.read(), f.name, 'exec'))"
-  "Template for Python session command with output results.
+  "Wrapper for session block with output results.
 
 Has a single %s escape, the tempfile containing the source code
 to evaluate.")
 
-(defconst org-babel-python--eval-ast "\
+(defconst org-babel-python--session-value-wrapper
+  (concat org-babel-python--def-format-value "
 import ast
-
 with open('%s') as f:
     __org_babel_python_ast = ast.parse(f.read())
 __org_babel_python_final = __org_babel_python_ast.body[-1]
-
 if isinstance(__org_babel_python_final, ast.Expr):
     __org_babel_python_ast.body = __org_babel_python_ast.body[:-1]
     exec(compile(__org_babel_python_ast, '<string>', 'exec'))
     __org_babel_python_final = eval(compile(ast.Expression(
         __org_babel_python_final.value), '<string>', 'eval'))
-    with open('%s', 'w') as f:
-        if %s:
-            import pprint
-            f.write(pprint.pformat(__org_babel_python_final))
-        else:
-            f.write(str(__org_babel_python_final))
 else:
     exec(compile(__org_babel_python_ast, '<string>', 'exec'))
-    __org_babel_python_final = None"
-  "Template for Python session command with value results.
+    __org_babel_python_final = None
+__org_babel_python_format_value(__org_babel_python_final, '%s', %s)")
+  "Wrapper for session block with value results.
 
 Has three %s escapes to be filled in:
 1. Tempfile containing source to evaluate.
 2. Tempfile to write results to.
-3. Whether to pretty print, \"True\" or \"False\".")
+3. result-params, converted from lisp to Python list.")
 
 (defun org-babel-python-evaluate
-  (session body &optional result-type result-params preamble)
+  (session body &optional result-type result-params preamble graphics-file)
   "Evaluate BODY as Python code."
   (if session
       (org-babel-python-evaluate-session
-       session body result-type result-params)
+       session body result-type result-params graphics-file)
     (org-babel-python-evaluate-external-process
-     body result-type result-params preamble)))
+     body result-type result-params preamble graphics-file)))
 
 (defun org-babel-python-evaluate-external-process
-    (body &optional result-type result-params preamble)
+    (body &optional result-type result-params preamble graphics-file)
   "Evaluate BODY in external python process.
 If RESULT-TYPE equals `output' then return standard output as a
 string.  If RESULT-TYPE equals `value' then return the value of the
@@ -294,16 +332,16 @@ (defun org-babel-python-evaluate-external-process
          (pcase result-type
            (`output (org-babel-eval org-babel-python-command
 				    (concat preamble (and preamble "\n")
-					    body)))
-           (`value (let ((tmp-file (org-babel-temp-file "python-")))
+					    (org-babel-python--output-graphics-wrapper
+					     body graphics-file))))
+           (`value (let ((results-file (or graphics-file
+				       (org-babel-temp-file "python-"))))
 		     (org-babel-eval
 		      org-babel-python-command
 		      (concat
 		       preamble (and preamble "\n")
 		       (format
-			(if (member "pp" result-params)
-			    org-babel-python-pp-wrapper-method
-			  org-babel-python-wrapper-method)
+			org-babel-python--nonsession-value-wrapper
 			(with-temp-buffer
 			  (python-mode)
 			  (insert body)
@@ -314,14 +352,15 @@ (defun org-babel-python-evaluate-external-process
 							 (line-end-position)))
 			   (forward-line 1))
 			 (buffer-string))
-			(org-babel-process-file-name tmp-file 'noquote))))
-		     (org-babel-eval-read-file tmp-file))))))
+			(org-babel-process-file-name results-file 'noquote)
+			(org-babel-python-var-to-python result-params))))
+		     (org-babel-eval-read-file results-file))))))
     (org-babel-result-cond result-params
       raw
       (org-babel-python-table-or-string (org-trim raw)))))
 
 (defun org-babel-python-evaluate-session
-    (session body &optional result-type result-params)
+    (session body &optional result-type result-params graphics-file)
   "Pass BODY to the Python process in SESSION.
 If RESULT-TYPE equals `output' then return standard output as a
 string.  If RESULT-TYPE equals `value' then return the value of the
@@ -334,17 +373,20 @@ (defun org-babel-python-evaluate-session
 	    (with-temp-file tmp-src-file (insert body))
 	    (pcase result-type
 	      (`output
-	       (let ((src-str (format org-babel-python--exec-tmpfile
-				      (org-babel-process-file-name
-				       tmp-src-file 'noquote))))
+	       (let ((src-str (org-babel-python--output-graphics-wrapper
+			       (format org-babel-python--session-output-wrapper
+				       (org-babel-process-file-name
+					tmp-src-file 'noquote))
+			       graphics-file)))
 		 (if (eq 'python-mode org-babel-python-mode)
 		     (py-send-string-no-output
 		      src-str (get-buffer-process session) session)
 		   (python-shell-send-string-no-output src-str))))
 	      (`value
-	       (let* ((results-file (org-babel-temp-file "python-"))
+	       (let* ((results-file (or graphics-file
+					(org-babel-temp-file "python-")))
 		      (src-str (format
-				org-babel-python--eval-ast
+				org-babel-python--session-value-wrapper
 				(org-babel-process-file-name tmp-src-file 'noquote)
 				(org-babel-process-file-name results-file 'noquote)
 				(org-babel-python-var-to-python result-params))))
-- 
2.28.0