bpo-31855: unittest.mock.mock_open() results now respects the argument of read([size]) (GH-11521)

unittest.mock.mock_open() results now respects the argument of read([size])

Co-Authored-By: remilapeyre <remi.lapeyre@henki.fr>

1 files changed, 14 insertions, 25 deletions

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index 351aba5..1e8057d 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()