25
25
__version__ = '1.0'
26
26
27
27
28
+ import io
28
29
import inspect
29
30
import pprint
30
31
import sys
@@ -2318,25 +2319,12 @@ def __init__(self, spec, spec_set=False, parent=None,
2318
2319
2319
2320
file_spec = None
2320
2321
2321
- def _iterate_read_data (read_data ):
2322
- # Helper for mock_open:
2323
- # Retrieve lines from read_data via a generator so that separate calls to
2324
- # readline, read, and readlines are properly interleaved
2325
- sep = b'\n ' if isinstance (read_data , bytes ) else '\n '
2326
- data_as_list = [l + sep for l in read_data .split (sep )]
2327
-
2328
- if data_as_list [- 1 ] == sep :
2329
10000
- # If the last line ended in a newline, the list comprehension will have an
2330
- # extra entry that's just a newline. Remove this.
2331
- data_as_list = data_as_list [:- 1 ]
2332
- else :
2333
- # If there wasn't an extra newline by itself, then the file being
2334
- # emulated doesn't have a newline to end the last line remove the
2335
- # newline that our naive format() added
2336
- data_as_list [- 1 ] = data_as_list [- 1 ][:- 1 ]
2337
2322
2338
- for line in data_as_list :
2339
- yield line
2323
+ def _to_stream (read_data ):
2324
+ if isinstance (read_data , bytes ):
2325
+ return io .BytesIO (read_data )
2326
+ else :
2327
+ return io .StringIO (read_data )
2340
2328
2341
2329
2342
2330
def mock_open (mock = None , read_data = '' ):
@@ -2351,20 +2339,23 @@ def mock_open(mock=None, read_data=''):
2351
2339
`read_data` is a string for the `read`, `readline` and `readlines` of the
2352
2340
file handle to return. This is an empty string by default.
2353
2341
"""
2342
+ _read_data = _to_stream (read_data )
2343
+ _state = [_read_data , None ]
2344
+
2354
2345
def _readlines_side_effect (* args , ** kwargs ):
2355
2346
if handle .readlines .return_value is not None :
2356
2347
return handle .readlines .return_value
2357
- return list ( _state [0 ])
2348
+ return _state [0 ]. readlines ( * args , ** kwargs )
2358
2349
2359
2350
def _read_side_effect (* args , ** kwargs ):
2360
2351
if handle .read .return_value is not None :
2361
2352
return handle .read .return_value
2362
- return type ( read_data )(). join ( _state [0 ])
2353
+ return _state [0 ]. read ( * args , ** kwargs )
2363
2354
2364
- def _readline_side_effect ():
2355
+ def _readline_side_effect (* args , ** kwargs ):
2365
2356
yield from _iter_side_effect ()
2366
2357
while True :
2367
- yield type ( read_data )( )
2358
+ yield _state [ 0 ]. readline ( * args , ** kwargs )
2368
2359
2369
2360
def _iter_side_effect ():
2370
2361
if handle .readline .return_value is not None :
@@ -2384,8 +2375,6 @@ def _iter_side_effect():
2384
2375
handle = MagicMock (spec = file_spec )
2385
2376
handle .__enter__ .return_value = handle
2386
2377
2387
- _state = [_iterate_read_data (read_data ), None ]
2388
-
2389
2378
handle .write .return_value = None
2390
2379
handle .read .return_value = None
2391
2380
handle .readline .return_value = None
@@ -2398,7 +2387,7 @@ def _iter_side_effect():
2398
2387
handle .__iter__ .side_effect = _iter_side_effect
2399
2388
2400
2389
def reset_data (* args , ** kwargs ):
2401
- _state [0 ] = _iterate_read_data (read_data )
2390
+ _state [0 ] = _to_stream (read_data )
2402
2391
if handle .readline .side_effect == _state [1 ]:
2403
2392
# Only reset the side effect if the user hasn't overridden it.
2404
2393
_state [1 ] = _readline_side_effect ()
0 commit comments