diff options
Diffstat (limited to 'Lib/unittest/mock.py')
| -rw-r--r-- | Lib/unittest/mock.py | 39 |
1 files changed, 14 insertions, 25 deletions
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index 351aba5d44..1e8057d5f5 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -25,6 +25,7 @@ __all__ = ( __version__ = '1.0' +import io import inspect import pprint import sys @@ -2318,25 +2319,12 @@ MethodWrapperTypes = ( file_spec = None -def _iterate_read_data(read_data): - # Helper for mock_open: - # Retrieve lines from read_data via a generator so that separate calls to - # readline, read, and readlines are properly interleaved - sep = b'\n' if isinstance(read_data, bytes) else '\n' - data_as_list = [l + sep for l in read_data.split(sep)] - - if data_as_list[-1] == sep: - # If the last line ended in a newline, the list comprehension will have an - # extra entry that's just a newline. Remove this. - data_as_list = data_as_list[:-1] - else: - # If there wasn't an extra newline by itself, then the file being - # emulated doesn't have a newline to end the last line remove the - # newline that our naive format() added - data_as_list[-1] = data_as_list[-1][:-1] - for line in data_as_list: - yield line +def _to_stream(read_data): + if isinstance(read_data, bytes): + return io.BytesIO(read_data) + else: + return io.StringIO(read_data) def mock_open(mock=None, read_data=''): @@ -2351,20 +2339,23 @@ def mock_open(mock=None, read_data=''): `read_data` is a string for the `read`, `readline` and `readlines` of the file handle to return. This is an empty string by default. """ + _read_data = _to_stream(read_data) + _state = [_read_data, None] + def _readlines_side_effect(*args, **kwargs): if handle.readlines.return_value is not None: return handle.readlines.return_value - return list(_state[0]) + return _state[0].readlines(*args, **kwargs) def _read_side_effect(*args, **kwargs): if handle.read.return_value is not None: return handle.read.return_value - return type(read_data)().join(_state[0]) + return _state[0].read(*args, **kwargs) - def _readline_side_effect(): + def _readline_side_effect(*args, **kwargs): yield from _iter_side_effect() while True: - yield type(read_data)() + yield _state[0].readline(*args, **kwargs) def _iter_side_effect(): if handle.readline.return_value is not None: @@ -2384,8 +2375,6 @@ def mock_open(mock=None, read_data=''): handle = MagicMock(spec=file_spec) handle.__enter__.return_value = handle - _state = [_iterate_read_data(read_data), None] - handle.write.return_value = None handle.read.return_value = None handle.readline.return_value = None @@ -2398,7 +2387,7 @@ def mock_open(mock=None, read_data=''): handle.__iter__.side_effect = _iter_side_effect def reset_data(*args, **kwargs): - _state[0] = _iterate_read_data(read_data) + _state[0] = _to_stream(read_data) if handle.readline.side_effect == _state[1]: # Only reset the side effect if the user hasn't overridden it. _state[1] = _readline_side_effect() |