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

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.

Lib/test/test_print.py

@@ -165,6 +165,23 @@ class TestPy2MigrationHint(unittest.TestCase):
 
         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:

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.

Objects/exceptions.c

@@ -2840,26 +2840,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);
+    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);