Skip to content

Commit

Permalink
capture: add type annotations to CaptureFixture
Browse files Browse the repository at this point in the history
It now has a str/bytes type parameter.
  • Loading branch information
bluetech committed Aug 10, 2020
1 parent 8a66f0a commit acc9310
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 30 deletions.
60 changes: 33 additions & 27 deletions src/_pytest/capture.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -495,32 +499,34 @@ 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 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:
Expand All @@ -543,7 +549,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

Expand All @@ -566,7 +572,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:
Expand Down Expand Up @@ -607,7 +613,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:
Expand All @@ -619,7 +625,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":
Expand Down Expand Up @@ -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 "<CaptureManager _method={!r} _global_capturing={!r} _capture_fixture={!r}>".format(
Expand Down Expand Up @@ -707,13 +713,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
Expand Down Expand Up @@ -812,14 +818,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

Expand All @@ -838,7 +844,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.
Expand Down Expand Up @@ -877,15 +883,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
Expand All @@ -894,15 +900,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
Expand All @@ -911,15 +917,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
Expand All @@ -928,15 +934,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
Expand Down
12 changes: 9 additions & 3 deletions testing/test_capture.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit acc9310

Please sign in to comment.