Jupyter notebook fix

mrjob/__init__.py

@@ -121,4 +121,4 @@
     'Andrea Zonca <andrea.zonca@gmail.com>',
 ]
 
-__version__ = '0.5.6'
+__version__ = '0.5.7.dev0'

mrjob/launch.py

@@ -37,7 +37,6 @@
 from mrjob.options import _alphabetize_options
 from mrjob.options import _fix_custom_options
 from mrjob.options import _print_help_for_groups
-from mrjob.py2 import PY2
 from mrjob.step import StepFailedException
 from mrjob.util import log_to_null
 from mrjob.util import log_to_stream
@@ -118,15 +117,15 @@ def error(msg):
         # Make it possible to redirect stdin, stdout, and stderr, for testing
         # See sandbox(), below.
         #
-        # These should always read/write bytes, not unicode
-        if PY2:
-            self.stdin = sys.stdin
-            self.stdout = sys.stdout
-            self.stderr = sys.stderr
-        else:
-            self.stdin = sys.stdin.buffer
-            self.stdout = sys.stdout.buffer
-            self.stderr = sys.stderr.buffer
+        # These should always read/write bytes, not unicode. Generally,
+        # on Python 2, sys.std* can read and write bytes, whereas on Python 3,
+        # you need to use sys.std*.buffer (which doesn't exist on Python 2).
+        #
+        # However, certain Python 3 environments, such as Jupyter notebook,
+        # act more like Python 2. See #1441.
+        self.stdin = getattr(sys.stdin, 'buffer', sys.stdin)
+        self.stdout = getattr(sys.stdout, 'buffer', sys.stdout)
+        self.stderr = getattr(sys.stderr, 'buffer', sys.stderr)
 
     @classmethod
     def _usage(cls):

tests/test_launch.py

@@ -26,6 +26,7 @@
 from mrjob.conf import combine_envs
 from mrjob.job import MRJob
 from mrjob.launch import MRJobLauncher
+from mrjob.py2 import PY2
 from mrjob.py2 import StringIO
 from mrjob.step import StepFailedException
 
@@ -352,3 +353,39 @@ def test_pass_through(self):
         self.assertEqual(self.get_value(MRRunner()), None)
         self.assertEqual(self.get_value(MRRunner(['-r', 'inline'])), 'inline')
         self.assertEqual(self.get_value(MRRunner(['-r', 'local'])), 'local')
+
+
+class StdStreamTestCase(TestCase):
+
+    def test_normal_python(self):
+        launcher = MRJobLauncher(args=['/path/to/script'])
+
+        if PY2:
+            self.assertEqual(launcher.stdin, sys.stdin)
+            self.assertEqual(launcher.stdout, sys.stdout)
+            self.assertEqual(launcher.stderr, sys.stderr)
+        else:
+            self.assertEqual(launcher.stdin, sys.stdin.buffer)
+            self.assertEqual(launcher.stdout, sys.stdout.buffer)
+            self.assertEqual(launcher.stderr, sys.stderr.buffer)
+
+    def test_python3_jupyter_notebook(self):
+        # regression test for #1441
+
+        # this actually works on any Python platform, since we use mocks
+        mock_stdin = Mock()
+        mock_stdin.buffer = Mock()
+
+        mock_stdout = Mock()
+        del mock_stdout.buffer
+
+        mock_stderr = Mock()
+        del mock_stderr.buffer
+
+        with patch.multiple(sys, stdin=mock_stdin,
+                            stdout=mock_stdout, stderr=mock_stderr):
+            launcher = MRJobLauncher(args=['/path/to/script'])
+
+        self.assertEqual(launcher.stdin, mock_stdin.buffer)
+        self.assertEqual(launcher.stdout, mock_stdout)
+        self.assertEqual(launcher.stderr, mock_stderr)