From acc9310c17eac6764beb36e9970ca84883369a2f Mon Sep 17 00:00:00 2001 From: Ran Benita Date: Sun, 9 Aug 2020 20:15:19 +0300 Subject: [PATCH] capture: add type annotations to CaptureFixture It now has a str/bytes type parameter. --- src/_pytest/capture.py | 60 ++++++++++++++++++++++------------------- testing/test_capture.py | 12 ++++++--- 2 files changed, 42 insertions(+), 30 deletions(-) diff --git a/src/_pytest/capture.py b/src/_pytest/capture.py index bf17e0eb8..7daf36ddd 100644 --- a/src/_pytest/capture.py +++ b/src/_pytest/capture.py @@ -6,7 +6,11 @@ import os 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,32 +499,34 @@ class FDCapture(FDCaptureBinary): # 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 in Python<3.5.3 due to https://bugs.python.org/issue31272 if sys.version_info >= (3, 5, 3): __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: @@ -543,7 +549,7 @@ class CaptureResult: return "CaptureResult(out={!r}, err={!r})".format(self.out, self.err) -class MultiCapture: +class MultiCapture(Generic[AnyStr]): _state = None _in_suspended = False @@ -566,7 +572,7 @@ class MultiCapture: 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: @@ -607,7 +613,7 @@ class MultiCapture: if self.in_: self.in_.done() - def readouterr(self) -> CaptureResult: + def readouterr(self) -> CaptureResult[AnyStr]: if self.out: out = self.out.snap() else: @@ -619,7 +625,7 @@ class MultiCapture: 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": @@ -657,8 +663,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 "".format( @@ -707,13 +713,13 @@ class CaptureManager: 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 @@ -812,14 +818,14 @@ class CaptureManager: 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 @@ -838,7 +844,7 @@ class CaptureFixture: 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. @@ -877,7 +883,7 @@ class CaptureFixture: @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 @@ -885,7 +891,7 @@ def capsys(request: SubRequest) -> Generator[CaptureFixture, None, None]: ``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 @@ -894,7 +900,7 @@ 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()`` @@ -902,7 +908,7 @@ def capsysbinary(request: SubRequest) -> Generator[CaptureFixture, None, None]: ``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 @@ -911,7 +917,7 @@ 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 @@ -919,7 +925,7 @@ def capfd(request: SubRequest) -> Generator[CaptureFixture, None, None]: ``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 @@ -928,7 +934,7 @@ 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 @@ -936,7 +942,7 @@ def capfdbinary(request: SubRequest) -> Generator[CaptureFixture, None, None]: ``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 diff --git a/testing/test_capture.py b/testing/test_capture.py index 92a0a3e41..b7545d73f 100644 --- a/testing/test_capture.py +++ b/testing/test_capture.py @@ -22,7 +22,9 @@ from _pytest.config import ExitCode # 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, @@ -30,7 +32,9 @@ def StdCaptureFD(out: bool = True, err: bool = True, in_: bool = True) -> MultiC ) -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, @@ -38,7 +42,9 @@ def StdCapture(out: bool = True, err: bool = True, in_: bool = True) -> MultiCap ) -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,