Source code for hbutils.testing.isolated.input

"""
Overview:
    Isolation for stdin stream.
"""
import io
import os
from contextlib import _RedirectStream, contextmanager
from typing import List, Union, ContextManager, TextIO

from hbutils.testing._base import TemporaryDirectory

__all__ = [
    'isolated_stdin',
]


# noinspection PyPep8Naming
class _redirect_stdin(_RedirectStream):
    _stream = 'stdin'


def _to_input_text(v: Union[str, List[str]]) -> str:
    if isinstance(v, str):
        return v
    elif isinstance(v, list):
        return '\n'.join(v)
    else:
        raise TypeError(f'Unknown stdin type - {v!r}.')


@contextmanager
def _stdin_via_mem(v: str) -> ContextManager[TextIO]:
    with io.StringIO(v) as f:
        yield f


@contextmanager
def _stdin_via_file(v: str) -> ContextManager[TextIO]:
    with TemporaryDirectory() as tdir:
        stdin_file = os.path.join(tdir, 'stdin')
        with open(stdin_file, 'w+') as inf:
            inf.write(v)

        with open(stdin_file, 'r') as f:
            yield f


[docs]@contextmanager def isolated_stdin(v: Union[str, List[str]], mem: bool = False): """ Overview: Isolation for stdin stream. :param v: Input content, a whole string or a list of string supported. :param mem: Use memory or not. Default is ``False`` which means \ a temporary file will be used as fake input stream. Examples:: >>> from hbutils.testing import isolated_stdin >>> with isolated_stdin(['123', '456']): ... a = int(input()) ... b = int(input()) ... print(a, b, a + b) 123 456 579 """ context = _stdin_via_mem if mem else _stdin_via_file with context(_to_input_text(v)) as f: with _redirect_stdin(f): yield