bpo-30721: Show correct syntax hint in Py3 when using Py2 redirection syntax (#2345)

1 files changed, 33 insertions, 0 deletions

diff --git a/Lib/test/test_print.py b/Lib/test/test_print.py
index 03f13b4..8099327 100644
--- a/Lib/test/test_print.py
+++ b/Lib/test/test_print.py
@@ -1,4 +1,5 @@
 import unittest
+import sys
 from io import StringIO
 
 from test import support
@@ -155,6 +156,38 @@ class TestPy2MigrationHint(unittest.TestCase):
 
         self.assertIn('print("Hello World", end=" ")', 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:
+            print >> sys.stderr, "message"
+        self.assertIn('Did you mean "print(<message>, '
+                'file=<output_stream>)', str(context.exception))
+
+        # Test correct hint is produced in the case where RHS implements
+        # __rrshift__ but returns NotImplemented
+        with self.assertRaises(TypeError) as context:
+            print >> 42
+        self.assertIn('Did you mean "print(<message>, '
+                'file=<output_stream>)', str(context.exception))
+
+        # Test stream redirection hint is specific to print
+        with self.assertRaises(TypeError) as context:
+            max >> sys.stderr
+        self.assertNotIn('Did you mean ', str(context.exception))
+
+        # Test stream redirection hint is specific to rshift
+        with self.assertRaises(TypeError) as context:
+            print << sys.stderr
+        self.assertNotIn('Did you mean', str(context.exception))
+
+        # Ensure right operand implementing rrshift still works
+        class OverrideRRShift:
+            def __rrshift__(self, lhs):
+                return 42 # Force result independent of LHS
+
+        self.assertEqual(print >> OverrideRRShift(), 42)
+
+
 
 if __name__ == "__main__":
     unittest.main()