[Python-checkins] bpo-32685: Improve suggestion for print statement (GH-5380)

[Nick Coghlan]

https://github.com/python/cpython/commit/b3b4b81d0147534151941105eba4af984acdc763
commit: b3b4b81d0147534151941105eba4af984acdc763
branch: 3.6
author: Miss Islington (bot) <31488909+miss-islington at users.noreply.github.com>
committer: Nick Coghlan <ncoghlan at gmail.com>
date: 2018-01-28T21:31:48+10:00
summary:

bpo-32685: Improve suggestion for print statement (GH-5380)

Better account for single-line compound statements and
semi-colon separated statements when suggesting
Py3 replacements for Py2 print statements.

Initial patch by Nitish Chandra.
(cherry picked from commit 43c0f1ac5ed8bc9c3bd048d2ce4de4c98a83de99)

files:
A Misc/NEWS.d/next/Core and Builtins/2018-01-28-12-25-06.bpo-32685.nGctze.rst
M Lib/test/test_print.py
M Objects/exceptions.c

diff --git a/Lib/test/test_print.py b/Lib/test/test_print.py
index 7bc23cf9fc63..e8381122182e 100644
--- a/Lib/test/test_print.py
+++ b/Lib/test/test_print.py
@@ -165,6 +165,23 @@ def test_string_with_leading_whitespace(self):
 
         self.assertIn('print("Hello World")', str(context.exception))
 
+    # bpo-32685: Suggestions for print statement should be proper when
+    # it is in the same line as the header of a compound statement
+    # and/or followed by a semicolon
+    def test_string_with_semicolon(self):
+        python2_print_str = 'print p;'
+        with self.assertRaises(SyntaxError) as context:
+            exec(python2_print_str)
+
+        self.assertIn('print(p)', str(context.exception))
+
+    def test_string_in_loop_on_same_line(self):
+        python2_print_str = 'for i in s: print i'
+        with self.assertRaises(SyntaxError) as context:
+            exec(python2_print_str)
+
+        self.assertIn('print(i)', str(context.exception))
+
     def test_stream_redirection_hint_for_py2_migration(self):
         # Test correct hint produced for Py2 redirection syntax
         with self.assertRaises(TypeError) as context:
diff --git a/Misc/NEWS.d/next/Core and Builtins/2018-01-28-12-25-06.bpo-32685.nGctze.rst b/Misc/NEWS.d/next/Core and Builtins/2018-01-28-12-25-06.bpo-32685.nGctze.rst
new file mode 100644
index 000000000000..07d7a072f6e1
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2018-01-28-12-25-06.bpo-32685.nGctze.rst	
@@ -0,0 +1,3 @@
+Improve suggestion when the Python 2 form of print statement is either
+present on the same line as the header of a compound statement or else
+terminated by a semi-colon instead of a newline. Patch by Nitish Chandra.
diff --git a/Objects/exceptions.c b/Objects/exceptions.c
index 0a75f1b0fe64..3bec50ea69e4 100644
--- a/Objects/exceptions.c
+++ b/Objects/exceptions.c
@@ -2859,26 +2859,31 @@ _PyErr_TrySetFromCause(const char *format, ...)
 static int
 _set_legacy_print_statement_msg(PySyntaxErrorObject *self, Py_ssize_t start)
 {
-    PyObject *strip_sep_obj = PyUnicode_FromString(" \t\r\n");
-    if (strip_sep_obj == NULL)
-        return -1;
-
-    // PRINT_OFFSET is to remove `print ` word from the data.
+    // PRINT_OFFSET is to remove the `print ` prefix from the data.
     const int PRINT_OFFSET = 6;
     const int STRIP_BOTH = 2;
-    // Issue 32028: Handle case when whitespace is used with print call
-    PyObject *initial_data = _PyUnicode_XStrip(self->text, STRIP_BOTH, strip_sep_obj);
-    if (initial_data == NULL) {
-        Py_DECREF(strip_sep_obj);
-        return -1;
+    Py_ssize_t start_pos = start + PRINT_OFFSET;
+    Py_ssize_t text_len = PyUnicode_GET_LENGTH(self->text);
+    Py_UCS4 semicolon = ';';
+    Py_ssize_t end_pos = PyUnicode_FindChar(self->text, semicolon,
+                                            start_pos, text_len, 1);
+    if (end_pos < -1) {
+      return -1;
+    } else if (end_pos == -1) {
+      end_pos = text_len;
     }
-    Py_ssize_t text_len = PyUnicode_GET_LENGTH(initial_data);
-    PyObject *data = PyUnicode_Substring(initial_data, PRINT_OFFSET, text_len);
-    Py_DECREF(initial_data);
+
+    PyObject *data = PyUnicode_Substring(self->text, start_pos, end_pos);
     if (data == NULL) {
-        Py_DECREF(strip_sep_obj);
         return -1;
     }
+
+    PyObject *strip_sep_obj = PyUnicode_FromString(" \t\r\n");
+    if (strip_sep_obj == NULL) {
+        Py_DECREF(data);
+        return -1;
+    }
+
     PyObject *new_data = _PyUnicode_XStrip(data, STRIP_BOTH, strip_sep_obj);
     Py_DECREF(data);
     Py_DECREF(strip_sep_obj);