From 1d532da49ec7e25ff9d78509a61c3aa82a29b482 Mon Sep 17 00:00:00 2001
From: Ran Benita <ran@unusedvar.com>
Date: Mon, 9 Nov 2020 15:07:51 +0200
Subject: [PATCH] assertion/rewrite: write pyc's according to PEP-552 on
 Python>=3.7

Python 3.7 changes the pyc format by adding a flags byte. Even though it
is not necessary for us to match it, it is nice to be able to read pyc
files we emit for debugging the rewriter.

Update our custom pyc files to use that format. We write flags==0
meaning we still use the mtime+size format rather the newer hash format.
---
 changelog/8014.trivial.rst       |  2 ++
 src/_pytest/assertion/rewrite.py | 34 +++++++++++++++------
 testing/test_assertrewrite.py    | 51 +++++++++++++++++++++++++++++++-
 3 files changed, 77 insertions(+), 10 deletions(-)
 create mode 100644 changelog/8014.trivial.rst

diff --git a/changelog/8014.trivial.rst b/changelog/8014.trivial.rst
new file mode 100644
index 000000000..3b9fb7bc2
--- /dev/null
+++ b/changelog/8014.trivial.rst
@@ -0,0 +1,2 @@
+`.pyc` files created by pytest's assertion rewriting now conform to the newer PEP-552 format on Python>=3.7.
+(These files are internal and only interpreted by pytest itself.)
diff --git a/src/_pytest/assertion/rewrite.py b/src/_pytest/assertion/rewrite.py
index 649726727..805d4c8b3 100644
--- a/src/_pytest/assertion/rewrite.py
+++ b/src/_pytest/assertion/rewrite.py
@@ -281,12 +281,16 @@ def _write_pyc_fp(
 ) -> None:
     # 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.
+    # import. However, there's little reason to deviate.
     fp.write(importlib.util.MAGIC_NUMBER)
+    # https://www.python.org/dev/peps/pep-0552/
+    if sys.version_info >= (3, 7):
+        flags = b"\x00\x00\x00\x00"
+        fp.write(flags)
     # 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
+    # "<LL" stands for 2 unsigned longs, little-endian.
     fp.write(struct.pack("<LL", mtime, size))
     fp.write(marshal.dumps(co))
 
@@ -365,21 +369,33 @@ def _read_pyc(
     except OSError:
         return None
     with fp:
+        # https://www.python.org/dev/peps/pep-0552/
+        has_flags = sys.version_info >= (3, 7)
         try:
             stat_result = os.stat(os.fspath(source))
             mtime = int(stat_result.st_mtime)
             size = stat_result.st_size
-            data = fp.read(12)
+            data = fp.read(16 if has_flags else 12)
         except OSError as e:
             trace(f"_read_pyc({source}): OSError {e}")
             return None
         # Check for invalid or out of date pyc file.
-        if (
-            len(data) != 12
-            or data[:4] != importlib.util.MAGIC_NUMBER
-            or struct.unpack("<LL", data[4:]) != (mtime & 0xFFFFFFFF, size & 0xFFFFFFFF)
-        ):
-            trace("_read_pyc(%s): invalid or out of date pyc" % source)
+        if len(data) != (16 if has_flags else 12):
+            trace("_read_pyc(%s): invalid pyc (too short)" % source)
+            return None
+        if data[:4] != importlib.util.MAGIC_NUMBER:
+            trace("_read_pyc(%s): invalid pyc (bad magic number)" % source)
+            return None
+        if has_flags and data[4:8] != b"\x00\x00\x00\x00":
+            trace("_read_pyc(%s): invalid pyc (unsupported flags)" % source)
+            return None
+        mtime_data = data[8 if has_flags else 4 : 12 if has_flags else 8]
+        if int.from_bytes(mtime_data, "little") != mtime & 0xFFFFFFFF:
+            trace("_read_pyc(%s): out of date" % source)
+            return None
+        size_data = data[12 if has_flags else 8 : 16 if has_flags else 12]
+        if int.from_bytes(size_data, "little") != size & 0xFFFFFFFF:
+            trace("_read_pyc(%s): invalid pyc (incorrect size)" % source)
             return None
         try:
             co = marshal.load(fp)
diff --git a/testing/test_assertrewrite.py b/testing/test_assertrewrite.py
index 09383cafe..84d5276e7 100644
--- a/testing/test_assertrewrite.py
+++ b/testing/test_assertrewrite.py
@@ -2,6 +2,7 @@ import ast
 import errno
 import glob
 import importlib
+import marshal
 import os
 import py_compile
 import stat
@@ -1063,12 +1064,60 @@ class TestAssertionRewriteHookDetails:
         py_compile.compile(str(source), str(pyc))
 
         contents = pyc.read_bytes()
-        strip_bytes = 20  # header is around 8 bytes, strip a little more
+        strip_bytes = 20  # header is around 16 bytes, strip a little more
         assert len(contents) > strip_bytes
         pyc.write_bytes(contents[:strip_bytes])
 
         assert _read_pyc(source, pyc) is None  # no error
 
+    @pytest.mark.skipif(
+        sys.version_info < (3, 7), reason="Only the Python 3.7 format for simplicity"
+    )
+    def test_read_pyc_more_invalid(self, tmp_path: Path) -> None:
+        from _pytest.assertion.rewrite import _read_pyc
+
+        source = tmp_path / "source.py"
+        pyc = tmp_path / "source.pyc"
+
+        source_bytes = b"def test(): pass\n"
+        source.write_bytes(source_bytes)
+
+        magic = importlib.util.MAGIC_NUMBER
+
+        flags = b"\x00\x00\x00\x00"
+
+        mtime = b"\x58\x3c\xb0\x5f"
+        mtime_int = int.from_bytes(mtime, "little")
+        os.utime(source, (mtime_int, mtime_int))
+
+        size = len(source_bytes).to_bytes(4, "little")
+
+        code = marshal.dumps(compile(source_bytes, str(source), "exec"))
+
+        # Good header.
+        pyc.write_bytes(magic + flags + mtime + size + code)
+        assert _read_pyc(source, pyc, print) is not None
+
+        # Too short.
+        pyc.write_bytes(magic + flags + mtime)
+        assert _read_pyc(source, pyc, print) is None
+
+        # Bad magic.
+        pyc.write_bytes(b"\x12\x34\x56\x78" + flags + mtime + size + code)
+        assert _read_pyc(source, pyc, print) is None
+
+        # Unsupported flags.
+        pyc.write_bytes(magic + b"\x00\xff\x00\x00" + mtime + size + code)
+        assert _read_pyc(source, pyc, print) is None
+
+        # Bad mtime.
+        pyc.write_bytes(magic + flags + b"\x58\x3d\xb0\x5f" + size + code)
+        assert _read_pyc(source, pyc, print) is None
+
+        # Bad size.
+        pyc.write_bytes(magic + flags + mtime + b"\x99\x00\x00\x00" + code)
+        assert _read_pyc(source, pyc, print) is None
+
     def test_reload_is_same_and_reloads(self, pytester: Pytester) -> None:
         """Reloading a (collected) module after change picks up the change."""
         pytester.makeini(