diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/lib/npyio.py | 16 | ||||
| -rw-r--r-- | numpy/lib/npyio.pyi | 2 | ||||
| -rw-r--r-- | numpy/lib/tests/test_io.py | 19 |
3 files changed, 35 insertions, 2 deletions
diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py index bfda6804a..f8f2ab7a2 100644 --- a/numpy/lib/npyio.py +++ b/numpy/lib/npyio.py @@ -167,6 +167,8 @@ class NpzFile(Mapping): >>> npz = np.load(outfile) >>> isinstance(npz, np.lib.npyio.NpzFile) True + >>> npz + NpzFile 'object' with keys x, y >>> sorted(npz.files) ['x', 'y'] >>> npz['x'] # getitem access @@ -178,6 +180,7 @@ class NpzFile(Mapping): # Make __exit__ safe if zipfile_factory raises an exception zip = None fid = None + _MAX_REPR_ARRAY_COUNT = 5 def __init__(self, fid, own_fid=False, allow_pickle=False, pickle_kwargs=None, *, @@ -259,6 +262,19 @@ class NpzFile(Mapping): else: raise KeyError("%s is not a file in the archive" % key) + def __repr__(self): + # Get filename or default to `object` + if isinstance(self.fid, str): + filename = self.fid + else: + filename = getattr(self.fid, "name", "object") + + # Get the name of arrays + array_names = ', '.join(self.files[:self._MAX_REPR_ARRAY_COUNT]) + if len(self.files) > self._MAX_REPR_ARRAY_COUNT: + array_names += "..." + return f"NpzFile {filename!r} with keys: {array_names}" + @set_module('numpy') def load(file, mmap_mode=None, allow_pickle=False, fix_imports=True, diff --git a/numpy/lib/npyio.pyi b/numpy/lib/npyio.pyi index 8007b2dc7..9dd3d6809 100644 --- a/numpy/lib/npyio.pyi +++ b/numpy/lib/npyio.pyi @@ -72,6 +72,7 @@ class NpzFile(Mapping[str, NDArray[Any]]): files: list[str] allow_pickle: bool pickle_kwargs: None | Mapping[str, Any] + _MAX_REPR_ARRAY_COUNT: int # Represent `f` as a mutable property so we can access the type of `self` @property def f(self: _T) -> BagObj[_T]: ... @@ -97,6 +98,7 @@ class NpzFile(Mapping[str, NDArray[Any]]): def __iter__(self) -> Iterator[str]: ... def __len__(self) -> int: ... def __getitem__(self, key: str) -> NDArray[Any]: ... + def __repr__(self) -> str: ... # NOTE: Returns a `NpzFile` if file is a zip file; # returns an `ndarray`/`memmap` otherwise diff --git a/numpy/lib/tests/test_io.py b/numpy/lib/tests/test_io.py index 06d6dbf8d..5a68fbc97 100644 --- a/numpy/lib/tests/test_io.py +++ b/numpy/lib/tests/test_io.py @@ -321,6 +321,21 @@ class TestSavezLoad(RoundtripTest): data.close() assert_(fp.closed) + @pytest.mark.parametrize("count, expected_repr", [ + (1, "NpzFile {fname!r} with keys: arr_0"), + (5, "NpzFile {fname!r} with keys: arr_0, arr_1, arr_2, arr_3, arr_4"), + # _MAX_REPR_ARRAY_COUNT is 5, so files with more than 5 keys are + # expected to end in '...' + (6, "NpzFile {fname!r} with keys: arr_0, arr_1, arr_2, arr_3, arr_4..."), + ]) + def test_repr_lists_keys(self, count, expected_repr): + a = np.array([[1, 2], [3, 4]], float) + with temppath(suffix='.npz') as tmp: + np.savez(tmp, *[a]*count) + l = np.load(tmp) + assert repr(l) == expected_repr.format(fname=tmp) + l.close() + class TestSaveTxt: def test_array(self): @@ -597,8 +612,8 @@ class TestSaveTxt: # in our process if needed, see gh-16889 memoryerror_raised = Value(c_bool) - # Since Python 3.8, the default start method for multiprocessing has - # been changed from 'fork' to 'spawn' on macOS, causing inconsistency + # Since Python 3.8, the default start method for multiprocessing has + # been changed from 'fork' to 'spawn' on macOS, causing inconsistency # on memory sharing model, lead to failed test for check_large_zip ctx = get_context('fork') p = ctx.Process(target=check_large_zip, args=(memoryerror_raised,)) |