From 45c4a8fb3d33698704ac17a376de36ada3cabecf Mon Sep 17 00:00:00 2001
From: Daniel Hahler <git@thequod.de>
Date: Thu, 7 Nov 2019 14:41:26 +0100
Subject: [PATCH] Use atomicrewrites only on Windows

Fixes https://github.com/pytest-dev/pytest/issues/6147
---
 changelog/6148.improvement.rst   |  1 +
 setup.py                         |  2 +-
 src/_pytest/assertion/rewrite.py | 69 +++++++++++++++++++++++---------
 testing/test_assertrewrite.py    | 34 ++++++++++------
 4 files changed, 74 insertions(+), 32 deletions(-)
 create mode 100644 changelog/6148.improvement.rst

diff --git a/changelog/6148.improvement.rst b/changelog/6148.improvement.rst
new file mode 100644
index 00000000000..c60d2c92142
--- /dev/null
+++ b/changelog/6148.improvement.rst
@@ -0,0 +1 @@
+``python-atomicwrites`` is only used on Windows, fixing a performance regression with assertion rewriting on Unix.
diff --git a/setup.py b/setup.py
index dcf63f6fd23..d7f3d7dbf61 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@
     "packaging",
     "attrs>=17.4.0",  # should match oldattrs tox env.
     "more-itertools>=4.0.0",
-    "atomicwrites>=1.0",
+    'atomicwrites>=1.0;sys_platform=="win32"',
     'pathlib2>=2.2.0;python_version<"3.6"',
     'colorama;sys_platform=="win32"',
     "pluggy>=0.12,<1.0",
diff --git a/src/_pytest/assertion/rewrite.py b/src/_pytest/assertion/rewrite.py
index 9c9d6135b1b..b8492993643 100644
--- a/src/_pytest/assertion/rewrite.py
+++ b/src/_pytest/assertion/rewrite.py
@@ -20,8 +20,6 @@
 from typing import Set
 from typing import Tuple
 
-import atomicwrites
-
 from _pytest._io.saferepr import saferepr
 from _pytest._version import version
 from _pytest.assertion import util
@@ -255,26 +253,59 @@ def get_data(self, pathname):
             return f.read()
 
 
-def _write_pyc(state, co, source_stat, pyc):
+def _write_pyc_fp(fp, source_stat, co):
     # Technically, we don't have to have the same pyc format as
     # (C)Python, since these "pycs" should never be seen by builtin
     # import. However, there's little reason deviate.
-    try:
-        with atomicwrites.atomic_write(fspath(pyc), mode="wb", overwrite=True) as fp:
-            fp.write(importlib.util.MAGIC_NUMBER)
-            # as of now, bytecode header expects 32-bit numbers for size and mtime (#4903)
-            mtime = int(source_stat.st_mtime) & 0xFFFFFFFF
-            size = source_stat.st_size & 0xFFFFFFFF
-            # "<LL" stands for 2 unsigned longs, little-ending
-            fp.write(struct.pack("<LL", mtime, size))
-            fp.write(marshal.dumps(co))
-    except EnvironmentError as e:
-        state.trace("error writing pyc file at {}: errno={}".format(pyc, e.errno))
-        # we ignore any failure to write the cache file
-        # there are many reasons, permission-denied, pycache dir being a
-        # file etc.
-        return False
-    return True
+    fp.write(importlib.util.MAGIC_NUMBER)
+    # as of now, bytecode header expects 32-bit numbers for size and mtime (#4903)
+    mtime = int(source_stat.st_mtime) & 0xFFFFFFFF
+    size = source_stat.st_size & 0xFFFFFFFF
+    # "<LL" stands for 2 unsigned longs, little-ending
+    fp.write(struct.pack("<LL", mtime, size))
+    fp.write(marshal.dumps(co))
+
+
+if sys.platform == "win32":
+    from atomicwrites import atomic_write
+
+    def _write_pyc(state, co, source_stat, pyc):
+        try:
+            with atomic_write(fspath(pyc), mode="wb", overwrite=True) as fp:
+                _write_pyc_fp(fp, source_stat, co)
+        except EnvironmentError as e:
+            state.trace("error writing pyc file at {}: errno={}".format(pyc, e.errno))
+            # we ignore any failure to write the cache file
+            # there are many reasons, permission-denied, pycache dir being a
+            # file etc.
+            return False
+        return True
+
+
+else:
+
+    def _write_pyc(state, co, source_stat, pyc):
+        proc_pyc = "{}.{}".format(pyc, os.getpid())
+        try:
+            fp = open(proc_pyc, "wb")
+        except EnvironmentError as e:
+            state.trace(
+                "error writing pyc file at {}: errno={}".format(proc_pyc, e.errno)
+            )
+            return False
+
+        try:
+            _write_pyc_fp(fp, source_stat, co)
+            os.rename(proc_pyc, fspath(pyc))
+        except BaseException as e:
+            state.trace("error writing pyc file at {}: errno={}".format(pyc, e.errno))
+            # we ignore any failure to write the cache file
+            # there are many reasons, permission-denied, pycache dir being a
+            # file etc.
+            return False
+        finally:
+            fp.close()
+        return True
 
 
 def _rewrite_test(fn, config):
diff --git a/testing/test_assertrewrite.py b/testing/test_assertrewrite.py
index 3aab3ac2afa..e2d6b89c8dd 100644
--- a/testing/test_assertrewrite.py
+++ b/testing/test_assertrewrite.py
@@ -959,24 +959,34 @@ def test_meta_path():
     def test_write_pyc(self, testdir, tmpdir, monkeypatch):
         from _pytest.assertion.rewrite import _write_pyc
         from _pytest.assertion import AssertionState
-        import atomicwrites
-        from contextlib import contextmanager
 
         config = testdir.parseconfig([])
         state = AssertionState(config, "rewrite")
-        source_path = tmpdir.ensure("source.py")
+        source_path = str(tmpdir.ensure("source.py"))
         pycpath = tmpdir.join("pyc").strpath
-        assert _write_pyc(state, [1], os.stat(source_path.strpath), pycpath)
+        assert _write_pyc(state, [1], os.stat(source_path), pycpath)
 
-        @contextmanager
-        def atomic_write_failed(fn, mode="r", overwrite=False):
-            e = IOError()
-            e.errno = 10
-            raise e
-            yield
+        if sys.platform == "win32":
+            from contextlib import contextmanager
 
-        monkeypatch.setattr(atomicwrites, "atomic_write", atomic_write_failed)
-        assert not _write_pyc(state, [1], source_path.stat(), pycpath)
+            @contextmanager
+            def atomic_write_failed(fn, mode="r", overwrite=False):
+                e = IOError()
+                e.errno = 10
+                raise e
+                yield
+
+            monkeypatch.setattr(
+                _pytest.assertion.rewrite, "atomic_write", atomic_write_failed
+            )
+        else:
+
+            def raise_ioerror(*args):
+                raise IOError()
+
+            monkeypatch.setattr("os.rename", raise_ioerror)
+
+        assert not _write_pyc(state, [1], os.stat(source_path), pycpath)
 
     def test_resources_provider_for_loader(self, testdir):
         """