Reintroduce memoryview() variant of copyfileobj()

...and use it from copyfile() only if WINDOWS and file size >= 128 MiB
python · giampaolo · Jun 19, 2018 · Jun 12, 2018 · Jun 13, 2018 · Jun 13, 2018
commit 8a3f88eb81560b520f13c9d48da779a262a49e18
diff --git a/Lib/shutil.py b/Lib/shutil.py
@@ -43,13 +43,14 @@
 except ImportError:
     getgrnam = None
 
+_WINDOWS = os.name == 'nt'
 posix = nt = None
 if os.name == 'posix':
     import posix
-elif os.name == 'nt':
+elif _WINDOWS:
     import nt
 
-COPY_BUFSIZE = 1024 * 1024 if os.name == 'nt' else 16 * 1024
+COPY_BUFSIZE = 1024 * 1024 if _WINDOWS else 16 * 1024
 _HAS_SENDFILE = posix and hasattr(os, "sendfile")
 _HAS_FCOPYFILE = posix and hasattr(posix, "_fcopyfile")  # macOS
 
@@ -168,6 +169,25 @@ def _fastcopy_sendfile(fsrc, fdst):
                 break  # EOF
             offset += sent
 
+def _copybinfileobj(fsrc, fdst, length=COPY_BUFSIZE):
+    """Copy 2 regular file objects open in binary mode.
+    This is used on Windows only for files >= 128 MiB as it appears to
+    give a considerable boost.
+    """
+    # Localize variable access to minimize overhead.
+    fsrc_readinto = fsrc.readinto
+    fdst_write = fdst.write
+    with memoryview(bytearray(length)) as mv:
+        while True:
+            n = fsrc_readinto(mv)
+            if not n:
+                break
+            elif n < length:
+                with mv[:n] as smv:
+                    fdst.write(smv)
+            else:
+                fdst_write(mv)
+
 def copyfileobj(fsrc, fdst, length=COPY_BUFSIZE):
     """copy data from file-like object fsrc to file-like object fdst"""
     # Localize variable access to minimize overhead.
@@ -192,7 +212,7 @@ def _samefile(src, dst):
             os.path.normcase(os.path.abspath(dst)))
 
 def copyfile(src, dst, *, follow_symlinks=True):
-    """Copy data from src to dst.
+    """Copy data from src to dst in the most efficient way possible.
 
     If follow_symlinks is not set and src is a symbolic link, a new
     symlink will be created instead of copying the file it points to.
@@ -201,7 +221,8 @@ def copyfile(src, dst, *, follow_symlinks=True):
     if _samefile(src, dst):
         raise SameFileError("{!r} and {!r} are the same file".format(src, dst))
 
-    for fn in [src, dst]:
+    fsize = 0
+    for i, fn in enumerate([src, dst]):
         try:
             st = os.stat(fn)
         except OSError:
@@ -211,6 +232,8 @@ def copyfile(src, dst, *, follow_symlinks=True):
             # XXX What about other special files? (sockets, devices...)
             if stat.S_ISFIFO(st.st_mode):
                 raise SpecialFileError("`%s` is a named pipe" % fn)
+            if _WINDOWS and i == 0:
+                fsize = st.st_size
 
     if not follow_symlinks and os.path.islink(src):
         os.symlink(os.readlink(src), dst)
@@ -230,7 +253,14 @@ def copyfile(src, dst, *, follow_symlinks=True):
                 except _GiveupOnFastCopy:
                     pass
 
-            copyfileobj(fsrc, fdst)
+            if _WINDOWS and fsize >= 128 * 1024 * 1024:
+                # Use alternate memoryview() based implementation on Windows
+                # for files >= 128 MiB. It appears this gives a considerable
+                # speedup, see:
+                # https://github.com/python/cpython/pull/7160#discussion_r195162475
+                _copybinfileobj(fsrc, fdst)
+            else:
+                copyfileobj(fsrc, fdst)
 
     return dst
 
@@ -1124,7 +1154,7 @@ def disk_usage(path):
         used = (st.f_blocks - st.f_bfree) * st.f_frsize
         return _ntuple_diskusage(total, used, free)
 
-elif os.name == 'nt':
+elif _WINDOWS:
 
     __all__.append('disk_usage')
     _ntuple_diskusage = collections.namedtuple('usage', 'total used free')

diff --git a/Lib/test/test_shutil.py b/Lib/test/test_shutil.py
@@ -1935,6 +1935,23 @@ def test_file_offset(self):
             self.assertEqual(src.tell(), self.FILESIZE)
             self.assertEqual(dst.tell(), self.FILESIZE)
 
+    @unittest.skipIf(os.name != 'nt', "Windows only")
+    def test_alternate_win_impl(self):
+        # On Windows copyfile() uses copyfileobj() for files < 128 MiB,
+        # else an alternate memoryview()-based implementation.
+        with unittest.mock.patch("shutil._copybinfileobj") as m:
+            shutil.copyfile(TESTFN, TESTFN2)
+        assert not m.called
+
+        fname = TESTFN + '-win'
+        self.addCleanup(support.unlink, fname)
+        write_test_file(fname, 128 * 1024 * 1024)
+        with unittest.mock.patch("shutil._copybinfileobj") as m:
+            shutil.copyfile(fname, TESTFN2)
+        assert m.called
+        shutil.copyfile(fname, TESTFN2)
+        self.assert_files_eq(fname, TESTFN2)
+
 
 class _ZeroCopyFileTest(object):
     """Tests common to all zero-copy APIs."""