capture: add type annotations to CaptureFixture

It now has a str/bytes type parameter.

Author: bluetech

src/_pytest/capture.py

@@ -6,7 +6,11 @@
 import sys
 from io import UnsupportedOperation
 from tempfile import TemporaryFile
+from typing import Any
+from typing import AnyStr
 from typing import Generator
+from typing import Generic
+from typing import Iterator
 from typing import Optional
 from typing import TextIO
 from typing import Tuple
@@ -495,31 +499,33 @@ def writeorg(self, data):
 # make it a namedtuple again.
 # [0]: https://github.com/python/mypy/issues/685
 @functools.total_ordering
-class CaptureResult:
+class CaptureResult(Generic[AnyStr]):
     """The result of :method:`CaptureFixture.readouterr`."""
 
     # Can't use slots due to bug in Python<=3.5.2: https://bugs.python.org/issue31272
     # __slots__ = ("out", "err")
 
-    def __init__(self, out, err) -> None:
-        self.out = out
-        self.err = err
+    def __init__(self, out: AnyStr, err: AnyStr) -> None:
+        self.out = out  # type: AnyStr
+        self.err = err  # type: AnyStr
 
     def __len__(self) -> int:
         return 2
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[AnyStr]:
         return iter((self.out, self.err))
 
-    def __getitem__(self, item: int):
+    def __getitem__(self, item: int) -> AnyStr:
         return tuple(self)[item]
 
-    def _replace(self, out=None, err=None) -> "CaptureResult":
+    def _replace(
+        self, *, out: Optional[AnyStr] = None, err: Optional[AnyStr] = None
+    ) -> "CaptureResult[AnyStr]":
         return CaptureResult(
             out=self.out if out is None else out, err=self.err if err is None else err
         )
 
-    def count(self, value) -> int:
+    def count(self, value: AnyStr) -> int:
         return tuple(self).count(value)
 
     def index(self, value) -> int:
@@ -542,7 +548,7 @@ def __repr__(self) -> str:
         return "CaptureResult(out={!r}, err={!r})".format(self.out, self.err)
 
 
-class MultiCapture:
+class MultiCapture(Generic[AnyStr]):
     _state = None
     _in_suspended = False
 
@@ -565,7 +571,7 @@ def start_capturing(self) -> None:
         if self.err:
             self.err.start()
 
-    def pop_outerr_to_orig(self):
+    def pop_outerr_to_orig(self) -> Tuple[AnyStr, AnyStr]:
         """Pop current snapshot out/err capture and flush to orig streams."""
         out, err = self.readouterr()
         if out:
@@ -606,7 +612,7 @@ def stop_capturing(self) -> None:
         if self.in_:
             self.in_.done()
 
-    def readouterr(self) -> CaptureResult:
+    def readouterr(self) -> CaptureResult[AnyStr]:
         if self.out:
             out = self.out.snap()
         else:
@@ -618,7 +624,7 @@ def readouterr(self) -> CaptureResult:
         return CaptureResult(out, err)
 
 
-def _get_multicapture(method: "_CaptureMethod") -> MultiCapture:
+def _get_multicapture(method: "_CaptureMethod") -> MultiCapture[str]:
     if method == "fd":
         return MultiCapture(in_=FDCapture(0), out=FDCapture(1), err=FDCapture(2))
     elif method == "sys":
@@ -656,8 +662,8 @@ class CaptureManager:
 
     def __init__(self, method: "_CaptureMethod") -> None:
         self._method = method
-        self._global_capturing = None  # type: Optional[MultiCapture]
-        self._capture_fixture = None  # type: Optional[CaptureFixture]
+        self._global_capturing = None  # type: Optional[MultiCapture[str]]
+        self._capture_fixture = None  # type: Optional[CaptureFixture[Any]]
 
     def __repr__(self) -> str:
         return "<CaptureManager _method={!r} _global_capturing={!r} _capture_fixture={!r}>".format(
@@ -706,13 +712,13 @@ def resume(self) -> None:
         self.resume_global_capture()
         self.resume_fixture()
 
-    def read_global_capture(self):
+    def read_global_capture(self) -> CaptureResult[str]:
         assert self._global_capturing is not None
         return self._global_capturing.readouterr()
 
     # Fixture Control
 
-    def set_fixture(self, capture_fixture: "CaptureFixture") -> None:
+    def set_fixture(self, capture_fixture: "CaptureFixture[Any]") -> None:
         if self._capture_fixture:
             current_fixture = self._capture_fixture.request.fixturename
             requested_fixture = capture_fixture.request.fixturename
@@ -811,14 +817,14 @@ def pytest_internalerror(self) -> None:
         self.stop_global_capturing()
 
 
-class CaptureFixture:
+class CaptureFixture(Generic[AnyStr]):
     """Object returned by the :py:func:`capsys`, :py:func:`capsysbinary`,
     :py:func:`capfd` and :py:func:`capfdbinary` fixtures."""
 
     def __init__(self, captureclass, request: SubRequest) -> None:
         self.captureclass = captureclass
         self.request = request
-        self._capture = None  # type: Optional[MultiCapture]
+        self._capture = None  # type: Optional[MultiCapture[AnyStr]]
         self._captured_out = self.captureclass.EMPTY_BUFFER
         self._captured_err = self.captureclass.EMPTY_BUFFER
 
@@ -837,7 +843,7 @@ def close(self) -> None:
             self._capture.stop_capturing()
             self._capture = None
 
-    def readouterr(self):
+    def readouterr(self) -> CaptureResult[AnyStr]:
         """Read and return the captured output so far, resetting the internal
         buffer.
 
@@ -876,15 +882,15 @@ def disabled(self) -> Generator[None, None, None]:
 
 
 @pytest.fixture
-def capsys(request: SubRequest) -> Generator[CaptureFixture, None, None]:
+def capsys(request: SubRequest) -> Generator[CaptureFixture[str], None, None]:
     """Enable text capturing of writes to ``sys.stdout`` and ``sys.stderr``.
 
     The captured output is made available via ``capsys.readouterr()`` method
     calls, which return a ``(out, err)`` namedtuple.
     ``out`` and ``err`` will be ``text`` objects.
     """
     capman = request.config.pluginmanager.getplugin("capturemanager")
-    capture_fixture = CaptureFixture(SysCapture, request)
+    capture_fixture = CaptureFixture[str](SysCapture, request)
     capman.set_fixture(capture_fixture)
     capture_fixture._start()
     yield capture_fixture
@@ -893,15 +899,15 @@ def capsys(request: SubRequest) -> Generator[CaptureFixture, None, None]:
 
 
 @pytest.fixture
-def capsysbinary(request: SubRequest) -> Generator[CaptureFixture, None, None]:
+def capsysbinary(request: SubRequest) -> Generator[CaptureFixture[bytes], None, None]:
     """Enable bytes capturing of writes to ``sys.stdout`` and ``sys.stderr``.
 
     The captured output is made available via ``capsysbinary.readouterr()``
     method calls, which return a ``(out, err)`` namedtuple.
     ``out`` and ``err`` will be ``bytes`` objects.
     """
     capman = request.config.pluginmanager.getplugin("capturemanager")
-    capture_fixture = CaptureFixture(SysCaptureBinary, request)
+    capture_fixture = CaptureFixture[bytes](SysCaptureBinary, request)
     capman.set_fixture(capture_fixture)
     capture_fixture._start()
     yield capture_fixture
@@ -910,15 +916,15 @@ def capsysbinary(request: SubRequest) -> Generator[CaptureFixture, None, None]:
 
 
 @pytest.fixture
-def capfd(request: SubRequest) -> Generator[CaptureFixture, None, None]:
+def capfd(request: SubRequest) -> Generator[CaptureFixture[str], None, None]:
     """Enable text capturing of writes to file descriptors ``1`` and ``2``.
 
     The captured output is made available via ``capfd.readouterr()`` method
     calls, which return a ``(out, err)`` namedtuple.
     ``out`` and ``err`` will be ``text`` objects.
     """
     capman = request.config.pluginmanager.getplugin("capturemanager")
-    capture_fixture = CaptureFixture(FDCapture, request)
+    capture_fixture = CaptureFixture[str](FDCapture, request)
     capman.set_fixture(capture_fixture)
     capture_fixture._start()
     yield capture_fixture
@@ -927,15 +933,15 @@ def capfd(request: SubRequest) -> Generator[CaptureFixture, None, None]:
 
 
 @pytest.fixture
-def capfdbinary(request: SubRequest) -> Generator[CaptureFixture, None, None]:
+def capfdbinary(request: SubRequest) -> Generator[CaptureFixture[bytes], None, None]:
     """Enable bytes capturing of writes to file descriptors ``1`` and ``2``.
 
     The captured output is made available via ``capfd.readouterr()`` method
     calls, which return a ``(out, err)`` namedtuple.
     ``out`` and ``err`` will be ``byte`` objects.
     """
     capman = request.config.pluginmanager.getplugin("capturemanager")
-    capture_fixture = CaptureFixture(FDCaptureBinary, request)
+    capture_fixture = CaptureFixture[bytes](FDCaptureBinary, request)
     capman.set_fixture(capture_fixture)
     capture_fixture._start()
     yield capture_fixture

testing/test_capture.py

@@ -22,23 +22,29 @@
 # pylib 1.4.20.dev2 (rev 13d9af95547e)
 
 
-def StdCaptureFD(out: bool = True, err: bool = True, in_: bool = True) -> MultiCapture:
+def StdCaptureFD(
+    out: bool = True, err: bool = True, in_: bool = True
+) -> MultiCapture[str]:
     return capture.MultiCapture(
         in_=capture.FDCapture(0) if in_ else None,
         out=capture.FDCapture(1) if out else None,
         err=capture.FDCapture(2) if err else None,
     )
 
 
-def StdCapture(out: bool = True, err: bool = True, in_: bool = True) -> MultiCapture:
+def StdCapture(
+    out: bool = True, err: bool = True, in_: bool = True
+) -> MultiCapture[str]:
     return capture.MultiCapture(
         in_=capture.SysCapture(0) if in_ else None,
         out=capture.SysCapture(1) if out else None,
         err=capture.SysCapture(2) if err else None,
     )
 
 
-def TeeStdCapture(out: bool = True, err: bool = True, in_: bool = True) -> MultiCapture:
+def TeeStdCapture(
+    out: bool = True, err: bool = True, in_: bool = True
+) -> MultiCapture[str]:
     return capture.MultiCapture(
         in_=capture.SysCapture(0, tee=True) if in_ else None,
         out=capture.SysCapture(1, tee=True) if out else None,