Re: [PATCH] Expanded ob-python results handling and plotting

[Jack Kamm <jackkamm@gmail.com>]

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

Hi Kyle,

Thanks for the comments, I'm attaching an updated patch.

Kyle Meyer <kyle@kyleam.com> writes:

> ModuleNotFoundError wasn't added until Python 3.6, so I think it'd be
> better to use its parent class, ImportError.

I did not know this, thanks for the tip.

> Should handling of Series also be added?

Yes, I've done so now. I'm not sure whether it's better to treat it like
a row or column vector, but since it has an "index", which are the row
names in a DataFrame, I decided to treat it as a column.


[-- 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: 12292 bytes --]

From 40db6b5497de78a9e69de219f4686b405db10c81 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      |  17 ++++++-
 lisp/ob-python.el | 126 +++++++++++++++++++++++++++++++---------------
 2 files changed, 100 insertions(+), 43 deletions(-)

diff --git a/etc/ORG-NEWS b/etc/ORG-NEWS
index 10658a970..75c945572 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,19 @@ 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/series, and will convert them to org-mode tables when
+appropriate.
+
+When the header argument =:results graphics= 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 before
+evaluation, and then plotted afterwards.
+
 *** =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..fb8fe380e 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,104 @@ (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 ImportError:
+                    pass
+                else:
+                    if isinstance(result, pd.DataFrame):
+                        result = [[''] + list(result.columns), None] + \
+[[i] + list(row) for i, row in result.iterrows()]
+                    elif isinstance(result, pd.Series):
+                        result = list(result.items())
+                try:
+                    import numpy as np
+                except ImportError:
+                    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 +334,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 +354,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,24 +375,27 @@ (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))))
 		 (if (eq 'python-mode org-babel-python-mode)
 		     (py-shell-send-string src-str (get-buffer-process session))
 		   (python-shell-send-string src-str))
-		 (sleep-for 0 5)
+		 (sleep-for 0 10)
 		 (org-babel-eval-read-file results-file)))))))
     (org-babel-result-cond result-params
       results
-- 
2.28.0