ob-python: Improvements to :return header argument

* lisp/ob-python.el (org-babel-execute:python): Allow return-val to be
non-nil in sessions, and concatenate it after the expanded body.

2 files changed, 60 insertions, 4 deletions

diff --git a/etc/ORG-NEWS b/etc/ORG-NEWS
index 6b580ee..5dc68cb 100644
--- a/etc/ORG-NEWS
+++ b/etc/ORG-NEWS
@@ -24,6 +24,59 @@ Earlier, IDs generated using =ts= method had a hard-coded format (i.e. =20200923
 The new option allows user to customise the format.
 Defaults are unchanged.
 
+** New features
+*** =ob-python= improvements to =:return= header argument 
+
+The =:return= header argument in =ob-python= now works for session
+blocks as well as non-session blocks.  Also, it now works with the
+=:epilogue= header argument -- previously, setting the =:return=
+header would cause the =:epilogue= to be ignored.
+
+This change allows more easily moving boilerplate out of the main code
+block and into the header. For example, for plotting, we need to add
+boilerplate to save the figure to a file and return the
+filename. Instead of doing this within the code block, we can now
+handle it through the header arguments as follows:
+
+#+BEGIN_SRC org
+,#+header: :var fname="/home/jack/tmp/plot.svg"
+,#+header: :epilogue plt.savefig(fname)
+,#+header: :return fname
+,#+begin_src python :results value file
+  import matplotlib, numpy
+  import matplotlib.pyplot as plt
+  fig=plt.figure(figsize=(4,2))
+  x=numpy.linspace(-15,15)
+  plt.plot(numpy.sin(x)/x)
+  fig.tight_layout()
+,#+end_src
+
+,#+RESULTS:
+[[file:/home/jack/tmp/plot.svg]]
+#+END_SRC
+
+As another example, we can use =:return= with the external [[https://pypi.org/project/tabulate/][tabulate]]
+package, to convert pandas Dataframes into orgmode tables:
+
+#+begin_src org
+,#+header: :prologue from tabulate import tabulate
+,#+header: :return tabulate(table, headers=table.columns, tablefmt="orgtbl")
+,#+begin_src python :results value raw :session
+  import pandas as pd
+  table = pd.DataFrame({
+      "a": [1,2,3],
+      "b": [4,5,6]
+  })
+,#+end_src
+
+,#+RESULTS:
+|   | a | b |
+|---+---+---|
+| 0 | 1 | 4 |
+| 1 | 2 | 5 |
+| 2 | 3 | 6 |
+#+end_src
+
 * Version 9.4
 ** Incompatible changes
 *** Possibly broken internal file links: please check and fix
diff --git a/lisp/ob-python.el b/lisp/ob-python.el
index 00a7c1a..785b919 100644
--- a/lisp/ob-python.el
+++ b/lisp/ob-python.el
@@ -81,13 +81,16 @@ This function is called by `org-babel-execute-src-block'."
 		   (cdr (assq :session 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))
+	 (return-val (when (eq result-type 'value)
 		       (cdr (assq :return params))))
 	 (preamble (cdr (assq :preamble params)))
          (full-body
-	  (org-babel-expand-body:generic
-	   (concat body (if return-val (format "\nreturn %s" return-val) ""))
-	   params (org-babel-variable-assignments:python params)))
+	  (concat
+	   (org-babel-expand-body:generic
+	    body params
+	    (org-babel-variable-assignments:python params))
+	   (when return-val
+	     (format (if session "\n%s" "\nreturn %s") return-val))))
          (result (org-babel-python-evaluate
 		  session full-body result-type result-params preamble)))
     (org-babel-reassemble-table