From 1fda86119055f7a5db8a30d867b26e802f6067ab Mon Sep 17 00:00:00 2001
From: Daniel Hahler <git@thequod.de>
Date: Sun, 8 Mar 2020 02:15:15 +0100
Subject: [PATCH] Fix crash when printing while capsysbinary is active

Previously, writing to sys.stdout/stderr in text-mode (e.g.
`print('foo')`) while a `capsysbinary` fixture is active, would crash
with:

    /usr/lib/python3.7/contextlib.py:119: in __exit__
        next(self.gen)
    E   TypeError: write() argument must be str, not bytes

This is due to some confusion in the types. The relevant functions are
`snap()` and `writeorg()`. The function `snap()` returns what was
captured, and the return type should be `bytes` for the binary captures
and `str` for the regular ones. The `snap()` return value is eventually
passed to `writeorg()` to be written to the original file, so it's input
type should correspond to `snap()`. But this was incorrect for
`SysCaptureBinary`, which handled it like `str`.

To fix this, be explicit in the `snap()` and `writeorg()`
implementations, also of the other Capture types.

We can't add type annotations yet, because the current inheritance
scheme breaks Liskov Substitution and mypy would complain. To be
refactored later.

Fixes: https://github.com/pytest-dev/pytest/issues/6871
Co-authored-by: Ran Benita (some modifications & commit message)
---
 changelog/6871.bugfix.rst |  1 +
 src/_pytest/capture.py    | 14 +++++++++++---
 testing/test_capture.py   | 36 +++++++++++++++++++++++++++++-------
 3 files changed, 41 insertions(+), 10 deletions(-)
 create mode 100644 changelog/6871.bugfix.rst

diff --git a/changelog/6871.bugfix.rst b/changelog/6871.bugfix.rst
new file mode 100644
index 00000000000..fe69c750915
--- /dev/null
+++ b/changelog/6871.bugfix.rst
@@ -0,0 +1 @@
+Fix crash with captured output when using the :fixture:`capsysbinary fixture <capsysbinary>`.
diff --git a/src/_pytest/capture.py b/src/_pytest/capture.py
index a64c72b5a7d..90de1d9fce7 100644
--- a/src/_pytest/capture.py
+++ b/src/_pytest/capture.py
@@ -570,8 +570,6 @@ def resume(self):
 
     def writeorg(self, data):
         """ write to original file descriptor. """
-        if isinstance(data, str):
-            data = data.encode("utf8")  # XXX use encoding of original stream
         os.write(self.targetfd_save, data)
 
 
@@ -591,6 +589,11 @@ def snap(self):
         self.tmpfile.truncate()
         return res
 
+    def writeorg(self, data):
+        """ write to original file descriptor. """
+        data = data.encode("utf-8")  # XXX use encoding of original stream
+        os.write(self.targetfd_save, data)
+
 
 class SysCaptureBinary:
 
@@ -642,8 +645,9 @@ def resume(self):
         self._state = "resumed"
 
     def writeorg(self, data):
-        self._old.write(data)
         self._old.flush()
+        self._old.buffer.write(data)
+        self._old.buffer.flush()
 
 
 class SysCapture(SysCaptureBinary):
@@ -655,6 +659,10 @@ def snap(self):
         self.tmpfile.truncate()
         return res
 
+    def writeorg(self, data):
+        self._old.write(data)
+        self._old.flush()
+
 
 class TeeSysCapture(SysCapture):
     def __init__(self, fd, tmpfile=None):
diff --git a/testing/test_capture.py b/testing/test_capture.py
index 65246151537..8b6e5967576 100644
--- a/testing/test_capture.py
+++ b/testing/test_capture.py
@@ -515,18 +515,40 @@ def test_hello(capfdbinary):
         reprec.assertoutcome(passed=1)
 
     def test_capsysbinary(self, testdir):
-        reprec = testdir.inline_runsource(
-            """\
+        p1 = testdir.makepyfile(
+            r"""
             def test_hello(capsysbinary):
                 import sys
-                # some likely un-decodable bytes
-                sys.stdout.buffer.write(b'\\xfe\\x98\\x20')
+
+                sys.stdout.buffer.write(b'hello')
+
+                # Some likely un-decodable bytes.
+                sys.stdout.buffer.write(b'\xfe\x98\x20')
+
+                sys.stdout.buffer.flush()
+
+                # Ensure writing in text mode still works and is captured.
+                # https://github.com/pytest-dev/pytest/issues/6871
+                print("world", flush=True)
+
                 out, err = capsysbinary.readouterr()
-                assert out == b'\\xfe\\x98\\x20'
+                assert out == b'hello\xfe\x98\x20world\n'
                 assert err == b''
+
+                print("stdout after")
+                print("stderr after", file=sys.stderr)
             """
         )
-        reprec.assertoutcome(passed=1)
+        result = testdir.runpytest(str(p1), "-rA")
+        result.stdout.fnmatch_lines(
+            [
+                "*- Captured stdout call -*",
+                "stdout after",
+                "*- Captured stderr call -*",
+                "stderr after",
+                "*= 1 passed in *",
+            ]
+        )
 
     def test_partial_setup_failure(self, testdir):
         p = testdir.makepyfile(
@@ -890,7 +912,7 @@ def test_writeorg(self, tmpfile):
         cap.start()
         tmpfile.write(data1)
         tmpfile.flush()
-        cap.writeorg(data2)
+        cap.writeorg(data2.decode("ascii"))
         scap = cap.snap()
         cap.done()
         assert scap == data1.decode("ascii")