Issue #4769: Fix main() function of the base64 module, use sys.stdin.buffer and

sys.stdout.buffer (instead of sys.stdin and sys.stdout) to use the bytes API

2 files changed, 36 insertions, 2 deletions

diff --git a/Lib/base64.py b/Lib/base64.py
index a8d9f21..20ed67d 100755
--- a/Lib/base64.py
+++ b/Lib/base64.py
@@ -383,9 +383,9 @@ def main():
         if o == '-u': func = decode
         if o == '-t': test(); return
     if args and args[0] != '-':
-        func(open(args[0], 'rb'), sys.stdout)
+        func(open(args[0], 'rb'), sys.stdout.buffer)
     else:
-        func(sys.stdin, sys.stdout)
+        func(sys.stdin.buffer, sys.stdout.buffer)
 
 
 def test():
diff --git a/Lib/test/test_base64.py b/Lib/test/test_base64.py
index fadee6d..7ff0db5 100644
--- a/Lib/test/test_base64.py
+++ b/Lib/test/test_base64.py
@@ -2,6 +2,8 @@ import unittest
 from test import support
 import base64
 import binascii
+import sys
+import subprocess
 
 
 
@@ -208,6 +210,38 @@ class BaseXYTestCase(unittest.TestCase):
 
 
 
+class TestMain(unittest.TestCase):
+    def get_output(self, *args, **options):
+        args = (sys.executable, '-m', 'base64') + args
+        return subprocess.check_output(args, **options)
+
+    def test_encode_decode(self):
+        output = self.get_output('-t')
+        self.assertSequenceEqual(output.splitlines(), (
+            b"b'Aladdin:open sesame'",
+            br"b'QWxhZGRpbjpvcGVuIHNlc2FtZQ==\n'",
+            b"b'Aladdin:open sesame'",
+        ))
+
+    def test_encode_file(self):
+        with open(support.TESTFN, 'wb') as fp:
+            fp.write(b'a\xffb\n')
+
+        output = self.get_output('-e', support.TESTFN)
+        self.assertEquals(output.rstrip(), b'Yf9iCg==')
+
+        with open(support.TESTFN, 'rb') as fp:
+            output = self.get_output('-e', stdin=fp)
+        self.assertEquals(output.rstrip(), b'Yf9iCg==')
+
+    def test_decode(self):
+        with open(support.TESTFN, 'wb') as fp:
+            fp.write(b'Yf9iCg==')
+        output = self.get_output('-d', support.TESTFN)
+        self.assertEquals(output, b'a\xffb\n')
+
+
+
 def test_main():
     support.run_unittest(__name__)