3 files changed, 85 insertions, 27 deletions

diff --git a/etc/ORG-NEWS b/etc/ORG-NEWS
index 6518c318d05..2068b3aabae 100644
--- a/etc/ORG-NEWS
+++ b/etc/ORG-NEWS
@@ -11,6 +11,15 @@ See the end of the file for license conditions.
 Please send Org bug reports to mailto:emacs-orgmode@gnu.org.
 
 * Version 9.4 (not yet released)
+** Incompatible changes
+*** 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,
+otherwise the result is None. Also, None will now show up under
+"#+RESULTS:", as it already did with ~:results value~ for non-session
+blocks.
+
 ** New features
 
 *** Numeric priorities are now allowed (up to 65)
diff --git a/lisp/ob-python.el b/lisp/ob-python.el
index 823f6e63d57..5f71577cb67 100644
--- a/lisp/ob-python.el
+++ b/lisp/ob-python.el
@@ -247,6 +247,25 @@ open('%s', 'w').write( pprint.pformat(main()) )")
    ")); "
    "__org_babel_python_fh.close()"))
 
+(defconst org-babel-python--eval-ast "\
+import ast
+try:
+    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'))
+    else:
+        exec(compile(__org_babel_python_ast, '<string>', 'exec'))
+        __org_babel_python_final = None
+except Exception:
+    from traceback import format_exc
+    __org_babel_python_final = format_exc()
+    raise")
+
 (defun org-babel-python-evaluate
   (session body &optional result-type result-params preamble)
   "Evaluate BODY as Python code."
@@ -294,32 +313,9 @@ 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))
                        (funcall send-wait)))
@@ -344,17 +340,35 @@ 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-ast
+                                   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))
+                   (dolist
+                       (statement
+                        (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)))))
+                     (insert statement)
+                     (funcall send-wait))
                    (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 48ca3d64088..7e2826404f2 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