ob-python: Fix :session :results output multiline behavior

* lisp/ob-python.el (orb-babel-python-evaluate-session): When :session
:results output, send multiline code blocks to tmpfile and execute in
Python with exec().
(org-babel-python--exec-tmpfile): New variable.
* testing/lisp/test-ob-python.el (test-ob-python/session-multiline):
Test for :session with multiple lines and indentation.

2 files changed, 43 insertions, 10 deletions

diff --git a/lisp/ob-python.el b/lisp/ob-python.el
index 60ec5fa..b3f50cf 100644
--- a/lisp/ob-python.el
+++ b/lisp/ob-python.el
@@ -239,6 +239,15 @@ def main():
 
 open('%s', 'w').write( pprint.pformat(main()) )")
 
+(defconst org-babel-python--exec-tmpfile
+  (concat
+   "__org_babel_python_fname = '%s'; "
+   "__org_babel_python_fh = open(__org_babel_python_fname); "
+   "exec(compile("
+   "__org_babel_python_fh.read(), __org_babel_python_fname, 'exec'"
+   ")); "
+   "__org_babel_python_fh.close()"))
+
 (defun org-babel-python-evaluate
   (session body &optional result-type result-params preamble)
   "Evaluate BODY as Python code."
@@ -306,16 +315,23 @@ last statement in BODY, as elisp."
          (results
           (pcase result-type
             (`output
-             (mapconcat
-              #'org-trim
-              (butlast
-               (org-babel-comint-with-output
-                   (session org-babel-python-eoe-indicator t body)
-                 (funcall input-body body)
-                 (funcall send-wait) (funcall send-wait)
-                 (insert org-babel-python-eoe-indicator)
-                 (funcall send-wait))
-               2) "\n"))
+	     (let ((body (if (string-match-p ".\n+." body) ; Multiline
+			     (let ((tmp-src-file (org-babel-temp-file
+						  "python-")))
+			       (with-temp-file tmp-src-file (insert body))
+			       (format org-babel-python--exec-tmpfile
+				       tmp-src-file))
+			   body)))
+	       (mapconcat
+		#'org-trim
+		(butlast
+		 (org-babel-comint-with-output
+		     (session org-babel-python-eoe-indicator t body)
+		   (funcall input-body body)
+		   (funcall send-wait) (funcall send-wait)
+		   (insert org-babel-python-eoe-indicator)
+		   (funcall send-wait))
+		 2) "\n")))
             (`value
              (let ((tmp-file (org-babel-temp-file "python-")))
                (org-babel-comint-with-output
diff --git a/testing/lisp/test-ob-python.el b/testing/lisp/test-ob-python.el
index d9fca22..915b1bc 100644
--- a/testing/lisp/test-ob-python.el
+++ b/testing/lisp/test-ob-python.el
@@ -101,6 +101,23 @@ return x
     (should (equal '(("col") ("a") ("b"))
 		   (org-babel-execute-src-block)))))
 
+(ert-deftest test-ob-python/session-multiline ()
+  ;; FIXME workaround to prevent starting prompt leaking into output
+  (run-python)
+  (sleep-for 0 10)
+  (org-test-with-temp-text "
+#+begin_src python :session :results output
+  foo = 0
+  for _ in range(10):
+      foo += 1
+
+      foo += 1
+
+  print(foo)
+#+end_src"
+   (org-babel-next-src-block)
+   (should (equal "20" (org-babel-execute-src-block)))))
+
 (provide 'test-ob-python)
 
 ;;; test-ob-python.el ends here