Changes preparing for stricter enforcement of bytes vs. str.

2 files changed, 15 insertions, 11 deletions

diff --git a/Lib/test/test_zipimport.py b/Lib/test/test_zipimport.py
index 17d5bb9..f38983f 100644
--- a/Lib/test/test_zipimport.py
+++ b/Lib/test/test_zipimport.py
@@ -255,7 +255,7 @@ class UncompressedZipImportTestCase(ImportHooksBaseTestCase):
         # stuff at the beginning of the file
         files = {TESTMOD + ".py": (NOW, test_src)}
         self.doTest(".py", files, TESTMOD,
-                    stuff="Some Stuff"*31)
+                    stuff=b"Some Stuff"*31)
 
     def assertModuleSource(self, module):
         self.assertEqual(inspect.getsource(module), test_src)
diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index e1fdc7f..f19ba31 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -909,10 +909,14 @@ class ZipFile:
         self.filelist.append(zinfo)
         self.NameToInfo[zinfo.filename] = zinfo
 
-    def writestr(self, zinfo_or_arcname, bytes):
-        """Write a file into the archive.  The contents is the string
-        'bytes'.  'zinfo_or_arcname' is either a ZipInfo instance or
+    def writestr(self, zinfo_or_arcname, data):
+        """Write a file into the archive.  The contents is 'data', which
+        may be either a 'str' or a 'bytes' instance; if it is a 'str',
+        it is encoded as UTF-8 first.
+        'zinfo_or_arcname' is either a ZipInfo instance or
         the name of the file in the archive."""
+        if isinstance(data, str):
+            data = data.encode("utf-8")
         if not isinstance(zinfo_or_arcname, ZipInfo):
             zinfo = ZipInfo(filename=zinfo_or_arcname,
                             date_time=time.localtime(time.time()))
@@ -924,21 +928,21 @@ class ZipFile:
             raise RuntimeError(
                   "Attempt to write to ZIP archive that was already closed")
 
-        zinfo.file_size = len(bytes)            # Uncompressed size
-        zinfo.header_offset = self.fp.tell()    # Start of header bytes
+        zinfo.file_size = len(data)            # Uncompressed size
+        zinfo.header_offset = self.fp.tell()    # Start of header data
         self._writecheck(zinfo)
         self._didModify = True
-        zinfo.CRC = binascii.crc32(bytes)       # CRC-32 checksum
+        zinfo.CRC = binascii.crc32(data)       # CRC-32 checksum
         if zinfo.compress_type == ZIP_DEFLATED:
             co = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
                  zlib.DEFLATED, -15)
-            bytes = co.compress(bytes) + co.flush()
-            zinfo.compress_size = len(bytes)    # Compressed size
+            data = co.compress(data) + co.flush()
+            zinfo.compress_size = len(data)    # Compressed size
         else:
             zinfo.compress_size = zinfo.file_size
-        zinfo.header_offset = self.fp.tell()    # Start of header bytes
+        zinfo.header_offset = self.fp.tell()    # Start of header data
         self.fp.write(zinfo.FileHeader())
-        self.fp.write(bytes)
+        self.fp.write(data)
         self.fp.flush()
         if zinfo.flag_bits & 0x08:
             # Write CRC and file sizes after the file data