Re: [PATCH] Fix several issues with python session value blocks

[Jack Kamm <jackkamm@gmail.com>]

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

Thanks Kyle! I've slightly updated the patch, because I had neglected to
handle Python exceptions. If you haven't started yet, please consider
the updated patch instead (it only differs by a few lines within
org-babel-python--eval-tmpfile).


[-- Warning: decoded text below may be mangled, UTF-8 assumed --]
[-- Attachment #2: 0001-ob-python-Fix-several-issues-with-session-results-va.patch --]
[-- Type: text/x-patch, Size: 6269 bytes --]

From c2fff65f36141b78d91af9d6264d0f936ee5a3a1 Mon Sep 17 00:00:00 2001
From: Jack Kamm <jackkamm@gmail.com>
Date: Mon, 20 Jan 2020 17:40:22 -0800
Subject: [PATCH] ob-python: Fix several issues with :session :results value

* lisp/ob-python.el (org-babel-python-evaluate-session): Fix a few
related issues with :session :results value blocks, including broken
if-else statements, indented blocks with blank lines, and returning
the wrong value when underscore has been used.  Uses the built-in ast
python module to parse a source block, execute it, and evaluate the
last line separately to return as a result.  Introduces a slight
change in behavior, requiring that the last line must be a top-level
statement if it's result is to be saved (otherwise, the result is
None).
---
 lisp/ob-python.el              | 68 +++++++++++++++++++---------------
 testing/lisp/test-ob-python.el | 35 +++++++++++++++++
 2 files changed, 73 insertions(+), 30 deletions(-)

diff --git a/lisp/ob-python.el b/lisp/ob-python.el
index 823f6e63d..0b1073df5 100644
--- a/lisp/ob-python.el
+++ b/lisp/ob-python.el
@@ -247,6 +247,24 @@ open('%s', 'w').write( pprint.pformat(main()) )")
    ")); "
    "__org_babel_python_fh.close()"))
 
+(defconst org-babel-python--eval-tmpfile "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]
+try:
+    if type(__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'))
+    else:
+        exec(compile(__org_babel_python_ast, '<string>', 'exec'))
+        __org_babel_python_final = None
+except Exception as e:
+        __org_babel_python_final = e
+        raise e")
+
 (defun org-babel-python-evaluate
   (session body &optional result-type result-params preamble)
   "Evaluate BODY as Python code."
@@ -294,34 +312,10 @@ If RESULT-TYPE equals `output' then return standard output as a
 string.  If RESULT-TYPE equals `value' then return the value of the
 last statement in BODY, as elisp."
   (let* ((send-wait (lambda () (comint-send-input nil t) (sleep-for 0 5)))
-	 (dump-last-value
-	  (lambda
-	    (tmp-file pp)
-	    (mapc
-	     (lambda (statement) (insert statement) (funcall send-wait))
-	     (if pp
-		 (list
-		  "import pprint"
-		  (format "open('%s', 'w').write(pprint.pformat(_))"
-			  (org-babel-process-file-name tmp-file 'noquote)))
-	       (list (format "open('%s', 'w').write(str(_))"
-			     (org-babel-process-file-name tmp-file
-                                                          'noquote)))))))
 	 (last-indent 0)
 	 (input-body (lambda (body)
-		       (dolist (line (split-string body "[\r\n]"))
-			 ;; Insert a blank line to end an indent
-			 ;; block.
-			 (let ((curr-indent (string-match "\\S-" line)))
-			   (if curr-indent
-			       (progn
-				 (when (< curr-indent last-indent)
-				   (insert "")
-				   (funcall send-wait))
-				 (setq last-indent curr-indent))
-			     (setq last-indent 0)))
-			 (insert line)
-			 (funcall send-wait))
+		       (mapc (lambda (line) (insert line) (funcall send-wait))
+			     (split-string body "[\r\n]"))
 		       (funcall send-wait)))
          (results
           (pcase result-type
@@ -344,17 +338,31 @@ last statement in BODY, as elisp."
 		   (funcall send-wait))
 		 2) "\n")))
             (`value
-             (let ((tmp-file (org-babel-temp-file "python-")))
+             (let ((tmp-results-file (org-babel-temp-file "python-"))
+		   (body (let ((tmp-src-file (org-babel-temp-file
+					      "python-")))
+			   (with-temp-file tmp-src-file (insert body))
+			   (format org-babel-python--eval-tmpfile
+				   tmp-src-file))))
                (org-babel-comint-with-output
                    (session org-babel-python-eoe-indicator nil body)
                  (let ((comint-process-echoes nil))
                    (funcall input-body body)
-                   (funcall dump-last-value tmp-file
-                            (member "pp" result-params))
+		   (mapc
+		    (lambda (statement) (insert statement) (funcall send-wait))
+		    (if (member "pp" result-params)
+			(list
+			 "import pprint"
+			 (format
+			  "open('%s', 'w').write(pprint.pformat(__org_babel_python_final))"
+			  (org-babel-process-file-name tmp-results-file 'noquote)))
+		      (list (format "open('%s', 'w').write(str(__org_babel_python_final))"
+				    (org-babel-process-file-name tmp-results-file
+								 'noquote)))))
                    (funcall send-wait) (funcall send-wait)
                    (insert org-babel-python-eoe-indicator)
                    (funcall send-wait)))
-               (org-babel-eval-read-file tmp-file))))))
+               (org-babel-eval-read-file tmp-results-file))))))
     (unless (string= (substring org-babel-python-eoe-indicator 1 -1) results)
       (org-babel-result-cond result-params
 	results
diff --git a/testing/lisp/test-ob-python.el b/testing/lisp/test-ob-python.el
index 48ca3d640..7e2826404 100644
--- a/testing/lisp/test-ob-python.el
+++ b/testing/lisp/test-ob-python.el
@@ -138,6 +138,41 @@ if True:
 	      (org-babel-execute-maybe)
 	      (org-babel-execute-src-block)))))
 
+(ert-deftest test-ob-python/if-else-block ()
+  (should
+   (equal "success" (org-test-with-temp-text "#+begin_src python :session :results value
+value = 'failure'
+if False:
+    pass
+else:
+    value = 'success'
+value
+#+end_src"
+	      (org-babel-execute-src-block)))))
+
+(ert-deftest test-ob-python/indent-block-with-blank-lines ()
+  (should
+   (equal 20
+	  (org-test-with-temp-text "#+begin_src python :session :results value
+  foo = 0
+  for i in range(10):
+      foo += 1
+
+      foo += 1
+
+  foo
+#+end_src"
+	    (org-babel-execute-src-block)))))
+
+(ert-deftest test-ob-python/assign-underscore ()
+  (should
+   (equal "success"
+	  (org-test-with-temp-text "#+begin_src python :session :results value
+_ = 'failure'
+'success'
+#+end_src"
+	    (org-babel-execute-src-block)))))
+
 (provide 'test-ob-python)
 
 ;;; test-ob-python.el ends here
-- 
2.25.0