diff options
author | Eric Wieser <wieser.eric@gmail.com> | 2018-02-08 21:40:21 -0800 |
---|---|---|
committer | Eric Wieser <wieser.eric@gmail.com> | 2018-09-23 09:39:53 -0700 |
commit | 8805080acf9e87a10881443c32b57d69fe44aec4 (patch) | |
tree | 3b24a4ef155bea93b1a1483d23460cd278dc490b /numpy/core/memmap.py | |
parent | 17a6221cb20e1f07b672d03bf69e3b526a11a03c (diff) | |
download | numpy-8805080acf9e87a10881443c32b57d69fe44aec4.tar.gz |
MAINT: Close the file if any unexpected errors occur within memmap
Diffstat (limited to 'numpy/core/memmap.py')
-rw-r--r-- | numpy/core/memmap.py | 130 |
1 files changed, 63 insertions, 67 deletions
diff --git a/numpy/core/memmap.py b/numpy/core/memmap.py index fd0a4bcd1..3cf101e43 100644 --- a/numpy/core/memmap.py +++ b/numpy/core/memmap.py @@ -2,7 +2,9 @@ from __future__ import division, absolute_import, print_function import numpy as np from .numeric import uint8, ndarray, dtype -from numpy.compat import long, basestring, is_pathlib_path +from numpy.compat import ( + long, basestring, is_pathlib_path, contextlib_nullcontext +) __all__ = ['memmap'] @@ -215,74 +217,68 @@ class memmap(ndarray): raise ValueError("shape must be given") if hasattr(filename, 'read'): - fid = filename - own_file = False + f_ctx = contextlib_nullcontext(filename) elif is_pathlib_path(filename): - fid = filename.open((mode == 'c' and 'r' or mode)+'b') - own_file = True + f_ctx = filename.open((mode == 'c' and 'r' or mode)+'b') else: - fid = open(filename, (mode == 'c' and 'r' or mode)+'b') - own_file = True - - fid.seek(0, 2) - flen = fid.tell() - descr = dtypedescr(dtype) - _dbytes = descr.itemsize - - if shape is None: - bytes = flen - offset - if bytes % _dbytes: - fid.close() - raise ValueError("Size of available data is not a " - "multiple of the data-type size.") - size = bytes // _dbytes - shape = (size,) - else: - if not isinstance(shape, tuple): - shape = (shape,) - size = np.intp(1) # avoid default choice of np.int_, which might overflow - for k in shape: - size *= k - - bytes = long(offset + size*_dbytes) - - if mode == 'w+' or (mode == 'r+' and flen < bytes): - fid.seek(bytes - 1, 0) - fid.write(b'\0') - fid.flush() - - if mode == 'c': - acc = mmap.ACCESS_COPY - elif mode == 'r': - acc = mmap.ACCESS_READ - else: - acc = mmap.ACCESS_WRITE - - start = offset - offset % mmap.ALLOCATIONGRANULARITY - bytes -= start - array_offset = offset - start - mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start) - - self = ndarray.__new__(subtype, shape, dtype=descr, buffer=mm, - offset=array_offset, order=order) - self._mmap = mm - self.offset = offset - self.mode = mode - - if isinstance(filename, basestring): - self.filename = os.path.abspath(filename) - elif is_pathlib_path(filename): - self.filename = filename.resolve() - # py3 returns int for TemporaryFile().name - elif (hasattr(filename, "name") and - isinstance(filename.name, basestring)): - self.filename = os.path.abspath(filename.name) - # same as memmap copies (e.g. memmap + 1) - else: - self.filename = None - - if own_file: - fid.close() + f_ctx = open(filename, (mode == 'c' and 'r' or mode)+'b') + + with f_ctx as fid: + fid.seek(0, 2) + flen = fid.tell() + descr = dtypedescr(dtype) + _dbytes = descr.itemsize + + if shape is None: + bytes = flen - offset + if bytes % _dbytes: + raise ValueError("Size of available data is not a " + "multiple of the data-type size.") + size = bytes // _dbytes + shape = (size,) + else: + if not isinstance(shape, tuple): + shape = (shape,) + size = np.intp(1) # avoid default choice of np.int_, which might overflow + for k in shape: + size *= k + + bytes = long(offset + size*_dbytes) + + if mode == 'w+' or (mode == 'r+' and flen < bytes): + fid.seek(bytes - 1, 0) + fid.write(b'\0') + fid.flush() + + if mode == 'c': + acc = mmap.ACCESS_COPY + elif mode == 'r': + acc = mmap.ACCESS_READ + else: + acc = mmap.ACCESS_WRITE + + start = offset - offset % mmap.ALLOCATIONGRANULARITY + bytes -= start + array_offset = offset - start + mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start) + + self = ndarray.__new__(subtype, shape, dtype=descr, buffer=mm, + offset=array_offset, order=order) + self._mmap = mm + self.offset = offset + self.mode = mode + + if isinstance(filename, basestring): + self.filename = os.path.abspath(filename) + elif is_pathlib_path(filename): + self.filename = filename.resolve() + # py3 returns int for TemporaryFile().name + elif (hasattr(filename, "name") and + isinstance(filename.name, basestring)): + self.filename = os.path.abspath(filename.name) + # same as memmap copies (e.g. memmap + 1) + else: + self.filename = None return self |