diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 067d20c..332db96 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -89,7 +89,7 @@ jobs: with: python-version: ${{ matrix.python }} - name: Install diskcache and test dependencies - run: pip install -e . "django==${{ matrix.django }}.*" pytest pytest-cov pytest-django pytest-xdist + run: pip install -e . "django==${{ matrix.django }}.*" numpy pytest pytest-cov pytest-django pytest-xdist - name: Run tests run: pytest @@ -134,7 +134,7 @@ jobs: sys.exit('this SQLite accepts double-quoted strings, so the job proves nothing') PY - name: Install diskcache and test dependencies - run: pip install -e . "django==5.2.*" pytest pytest-cov pytest-django pytest-xdist + run: pip install -e . "django==5.2.*" numpy pytest pytest-cov pytest-django pytest-xdist - name: Run tests run: pytest diff --git a/.pylintrc b/.pylintrc index 328db8f..eb4a955 100644 --- a/.pylintrc +++ b/.pylintrc @@ -341,7 +341,7 @@ indent-string=' ' max-line-length=100 # Maximum number of lines in a module. -max-module-lines=4000 +max-module-lines=4100 # Allow the body of a class to be on the same line as the declaration if body # contains single statement. diff --git a/CHANGELOG.md b/CHANGELOG.md index 9d2a2f8..61ba1f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,8 +8,12 @@ diskcache-community forked from DiskCache at upstream commit ebfa37c (3 March 20 - `SafeJSONDisk` stores JSON like `JSONDisk`, reads and writes the same files, and never reads a pickle. `JSONDisk` writes only JSON but inherits the pickle reads of `Disk`, so a crafted row still ran code and a crafted path could still delete a file outside the cache. `SafeJSONDisk` refuses pickled keys and values without parsing them and confines paths as `SafeDisk` does (PR #33). +- `SafeDisk.with_numpy()` returns a `SafeDisk` subclass that stores numpy arrays in the `.npy` format instead of pickling them, through pickle's persistent IDs, anywhere in a value including inside classes allowed with `with_types`, which combines with it in either order, as `with_options` does. numpy scalars read back as the same type. Reading a large array takes no second copy of it. Headers are checked against the data, so forged arrays are refused rather than raising `MemoryError`, and every forged entry raises `UnpicklingError` or reads as a miss with `on_error='miss'`. Object arrays, `ndarray` subclasses and arrays the reader would refuse raise `TypeError` at write time. Entries holding numpy data can't be read by other disks, including a plain `SafeDisk`, or by 5.6.3, which raise `UnpicklingError`. numpy stays optional, and `SafeDisk` without `with_numpy()` still refuses it; its `with_types` error now points to `with_numpy()` (PR #34). +- `DISKCACHE_SAFE_TYPES` and `DISKCACHE_SAFE_MODULES` now apply when `DISKCACHE_DISK` names a `SafeDisk` subclass, not only with `safe` (PR #34). + ### Fixes +- A float NaN stored as a value is kept, on every disk. It was written to SQLite, which stores NaN as NULL, so it read back as `None`. It is now stored as a pickle, which 5.6.3 reads too (PR #34). - `memoize_stampede` works with `SafeDisk`: the key of its early recomputation marker held the internal `ENOVAL` constant, which `SafeDisk` refuses, so a call raised `TypeError` once a cached value neared expiry (PR #32). ## 5.8.1 (2026-10-18) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e6c2134..fe131ff 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -36,16 +36,3 @@ tox -e bluecheck,isortcheck,flake8,mypy,pylint,doc8,rstcheck,docs Work on a branch and open a pull request against `master`. CI runs the lint checks, the test suite on Python 3.10 to 3.14 on Linux, macOS and Windows, the Django 4.2 and 6.0 jobs and an import check without Django, and a pull request merges once CI is green and the change has been reviewed. Keep each pull request to one change, and write commit subjects as one sentence that says what the change does. In the description, say what was wrong, list the concrete changes, and say what you ran to check them. Fixes imported from upstream pull requests keep their original commits and authors; the description says which upstream pull request is imported. Entries for merged changes go into [CHANGELOG.md](CHANGELOG.md) under `Unreleased`. - -## Compatibility contract - -This fork publishes the same library under a new distribution name, `diskcache-community`, and a new import name, `diskcache_community`, so it installs next to the upstream `diskcache` package. The following stay as upstream left them, and a change to any of them is a compatibility break that needs a major version: - -- the import name, `import diskcache_community` (it replaced `import diskcache` in 5.8.0, before any dependent used the fork), and the public API exported by `diskcache_community/__init__.py`; -- the `cache.db` schema: the `Settings` and `Cache` tables, the value modes 1 to 4, and the index and trigger names; -- the directory layout: `/ab/cd/<28 hex>.val` for values and `%03d` for `FanoutCache` shards; -- the `RAW`, `BINARY`, `TEXT` and `PICKLE` serialization and the bytes `JSONDisk` writes; -- the key format `memoize` derives from a call; -- the persisted `Settings` keys; no new `disk_*` or `Settings` key before 6.0.0. - -A process running 5.6.3 must be able to open a directory written by this fork, and the reverse. Golden fixtures written by the 5.6.3 wheel guard this in the test suite. diff --git a/README.md b/README.md index 8be58df..e7b34ee 100644 --- a/README.md +++ b/README.md @@ -112,8 +112,8 @@ cache. semaphores and throttling - Thread-safe and process-safe, with fork-safe connection handling - Read-only mode, automatic when the directory is not writable -- Opt-in `SafeDisk` and `SafeJSONDisk` for caches whose directory other users - can write to +- Opt-in `SafeDisk` and `SafeJSONDisk` for caches whose directory other + users can write to, with `SafeDisk.with_numpy()` for numpy arrays - Pure Python on top of the standard library `sqlite3`; no server, no compiler - Tested on CPython 3.10 to 3.14 on Linux, macOS and Windows - Gets in about 2.6 µs and sets in about 19 µs on an Apple M5 Pro in the diff --git a/SECURITY.md b/SECURITY.md index 1303971..aa292dc 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -19,12 +19,9 @@ Where this fork stands: - 5.7.0, the first release of this fork, reads pickles with the default `Disk` exactly as 5.6.3 does, so it opens every 5.6.3 cache. It adds an opt-in `SafeDisk` that unpickles only an allowlist of standard-library types, rejects anything outside it at write time and confines value paths to the cache directory; `Cache`, `FanoutCache`, `Deque` and `Index` take it through `disk=SafeDisk` and the Django backend through the `DISK` option. New cache directories are created with mode `0700`. - 5.8.0 is imported as `diskcache_community`, next to the upstream `diskcache` package. `SafeDisk.with_types` allows your own classes and the classes defined in named modules while everything else stays refused; numpy classes cannot be allowed, because numpy unpickling is not memory-safe. `configure()` and `DISKCACHE_DISK=safe` select `SafeDisk` for caches created without a disk, including those that libraries create through `diskcache_community`. `Disk.with_options(on_error='miss')` deletes entries that cannot be read and treats them as missing. The steps are in [Fix CVE-2025-69872 in diskcache](https://diskcache-community.readthedocs.io/en/latest/fix-cve-2025-69872.html). - `JSONDisk` writes no pickles but still reads one from a crafted row, as it does in 5.6.3. `SafeJSONDisk` reads and writes the same files and never reads a pickle; `SafeDisk` does not decode what `JSONDisk` wrote. +- `SafeDisk.with_numpy()` returns a `SafeDisk` that also stores numpy arrays, in the `.npy` format rather than as pickles. numpy pickles stay refused, because unpickling an array is not memory-safe. - The default stays as it is. Existing caches hold objects outside any allowlist, such as instances of your own classes, Django responses and the keys of `memoize(typed=True)`, so `SafeDisk` remains an explicit choice. -Unless you have opted into `SafeDisk` or `SafeJSONDisk`, treat the cache directory as trusted input: keep it on local storage that only the user running the process can write to, and do not share it with other accounts. - -A note on scanners: vulnerability scanners match on the distribution name, so `pip-audit`, Dependabot and similar tools report nothing for `diskcache-community` even though the default `Disk` behaves as it did in 5.6.3. A clean scan is not evidence that this issue is fixed. Use the list above instead, and the release notes of the version you run. - ## Dependencies The library has no runtime dependencies. Advisories for the development and CI dependencies are tracked through Dependabot alerts on this repository. diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py new file mode 100644 index 0000000..f3beab4 --- /dev/null +++ b/diskcache_community/arrays.py @@ -0,0 +1,325 @@ +"""Disk that stores numpy arrays without pickling them.""" + +import io +import math +import pickle +import warnings + +from .core import SAFE_PICKLE_NAMES, Disk, SafeUnpickler, _ReaderMissing + +# Pickle protocol 5 is the first that stores a buffer in place, so that +# writing an array copies it once, into the row, and reading a file-backed +# entry needs no second copy of it. +_MIN_PROTOCOL = 5 + +# The numpy scalar types that are stored, by the name written to the row, with +# the Python type that ``item()`` gives for each and the reader insists on. The +# names do not vary by platform or numpy version, unlike the types' own names. +_SCALARS = { + 'bool': ('bool_', bool), + 'int8': ('int8', int), + 'int16': ('int16', int), + 'int32': ('int32', int), + 'int64': ('int64', int), + 'uint8': ('uint8', int), + 'uint16': ('uint16', int), + 'uint32': ('uint32', int), + 'uint64': ('uint64', int), + 'float16': ('float16', float), + 'float32': ('float32', float), + 'float64': ('float64', float), + 'complex64': ('complex64', complex), + 'complex128': ('complex128', complex), + 'str': ('str_', str), + 'bytes': ('bytes_', bytes), +} + + +def _import_numpy(): + # pylint: disable=import-error,import-outside-toplevel + try: + import numpy + except ImportError as error: + raise _ReaderMissing( + 'SafeDisk needs numpy to read this entry' + ) from error + return numpy + + +def _read_header(numpy, header, size): + """Return shape, Fortran order and dtype from the ``.npy`` `header`. + + Raise UnpicklingError unless the header is consistent with `size` bytes + of data. The sizes are checked because a forged shape must not decide + what gets read. Anything that goes wrong while parsing, including a + MemoryError from the Python parser on a crafted header, which + ``on_error='miss'`` would not treat as an unreadable entry, is reported + as UnpicklingError. + + """ + # pylint: disable=broad-except,no-member + try: + reader = io.BytesIO(header) + version = numpy.lib.format.read_magic(reader) + if version == (1, 0): + fields = numpy.lib.format.read_array_header_1_0(reader) + elif version == (2, 0): + fields = numpy.lib.format.read_array_header_2_0(reader) + else: + raise ValueError('unsupported .npy version %r' % (version,)) + except Exception as error: + raise pickle.UnpicklingError( + 'SafeDisk refuses the array header' + ) from error + + shape, fortran, dtype = fields + count = math.prod(shape) + + # A zero-width dtype would let any shape match any amount of data. + if ( + dtype.hasobject + or reader.tell() != len(header) + or dtype.itemsize * count != size + or (count and not dtype.itemsize) + ): + raise pickle.UnpicklingError('SafeDisk refuses the array header') + + return shape, fortran, dtype + + +def _pack_array(numpy, array): + """Return the ``.npy`` header of `array` and a buffer of its data. + + The buffer is the array's own memory, so pickle copies it once, into + the row, as a bytearray. An array that is not contiguous or not + writable is copied first, since pickle stores a read-only buffer as + bytes, which the reader refuses. + + """ + # pylint: disable=no-member + flags = array.flags + if not (flags.writeable and (flags.c_contiguous or flags.f_contiguous)): + array = array.copy(order='A') + + buffer = io.BytesIO() + try: + with warnings.catch_warnings(): + # numpy warns about dtypes it would save as pickles or in a + # format the reader refuses; both are refused below. + warnings.simplefilter('ignore') + # Version 1 holds every header the reader accepts: numpy refuses + # a header over 10000 bytes, and version 1 holds 65535. + fields = numpy.lib.format.header_data_from_array_1_0(array) + numpy.lib.format.write_array_header_1_0(buffer, fields) + header = buffer.getvalue() + _read_header(numpy, header, array.nbytes) + except (ValueError, pickle.UnpicklingError) as error: + raise TypeError( + 'SafeDisk cannot store an array of dtype %s' % (array.dtype,) + ) from error + + return header, pickle.PickleBuffer(array) + + +def _unpack_array(header, data): + """Return the array in the bytearray `data`, which it keeps as its data. + + A view takes no second copy of the array, and stays writable because + `data` is. `header` and `data` may have been tampered with. + + """ + # pylint: disable=broad-except + numpy = _import_numpy() + shape, fortran, dtype = _read_header(numpy, header, len(data)) + + try: + if not math.prod(shape): + return numpy.empty(shape, dtype) + array = numpy.frombuffer(data, dtype) + return array.reshape(shape, order='F' if fortran else 'C') + except Exception as error: + raise pickle.UnpicklingError( + 'SafeDisk refuses the array data' + ) from error + + +def _scalar_name(scalar): + """Return the name to store numpy `scalar` under, or raise TypeError.""" + dtype = scalar.dtype + name = {'U': 'str', 'S': 'bytes'}.get(dtype.kind, dtype.name) + if name not in _SCALARS: + raise TypeError( + 'SafeDisk cannot store a scalar of type numpy.%s' + % (type(scalar).__name__,) + ) + return name + + +def _unpack_scalar(name, value): + """Return the numpy scalar `value` of type `name`, which may be forged. + + The value must be of the Python type the name is built from: numpy + would otherwise build something else, such as ``bytes_(n)``, which is + `n` zero bytes. + + """ + # pylint: disable=broad-except,unidiomatic-typecheck + entry = _SCALARS.get(name) if type(name) is str else None + if entry is None or type(value) is not entry[1]: + raise pickle.UnpicklingError('SafeDisk refuses a scalar') + + numpy = _import_numpy() + + try: + return getattr(numpy, entry[0])(value) + except Exception as error: + raise pickle.UnpicklingError('SafeDisk refuses a scalar') from error + + +class _ArrayPickler(pickle.Pickler): + """Pickler that writes arrays and numpy scalars as persistent IDs.""" + + def __init__(self, file, protocol, numpy): + super().__init__(file, protocol=protocol) + self._numpy = numpy + self._arrays = {} + self._scalar_names = {} + self.reduced = False + + def reducer_override(self, obj): + """Note that an object was reduced, which can look up a name. + + CPython skips this hook only for None, booleans and exact ints, + floats, bytes, strings, dicts, sets, frozensets, lists and tuples, + none of which pickle a name, so every pickle that looks one up is + noted here and scanned. + + """ + # pylint: disable=unused-argument + self.reduced = True + return NotImplemented + + def persistent_id(self, obj): + # pylint: disable=unidiomatic-typecheck + numpy = self._numpy + if isinstance(obj, numpy.ndarray): + if type(obj) is not numpy.ndarray: + raise TypeError( + 'SafeDisk cannot store a subclass of ndarray: %s.%s' + % (type(obj).__module__, type(obj).__qualname__) + ) + # pickle's own memo skips persistent IDs, and a value may hand it + # temporary arrays. The buffer in the ID keeps `obj` alive, so + # its id is not reused while this memo holds the ID. + pid = self._arrays.get(id(obj)) + if pid is None: + pid = self._arrays[id(obj)] = ('npy',) + _pack_array( + numpy, obj + ) + return pid + if isinstance(obj, numpy.generic): + kind = type(obj) + name = self._scalar_names.get(kind) + if name is None: + name = self._scalar_names[kind] = _scalar_name(obj) + value = obj.item() + if isinstance(value, numpy.generic): + raise TypeError( + 'SafeDisk cannot store a scalar of type numpy.%s' + % (kind.__name__,) + ) + return ('scalar', name, value) + return None + + +class _ArrayUnpickler(SafeUnpickler): + """Unpickler that reads back what _ArrayPickler wrote. + + A forged row can follow an array with BUILD, which calls its + ``__setstate__``. numpy restores an array's internal flags from the dtype + in that state without checking them, which is why numpy pickles are + refused. The call is harmless here because the dtype slot accepts only a + ``numpy.dtype``, and neither the allowlist nor a persistent ID can make + one; a test pins that. + + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._arrays = {} + self._strings = {} + + def persistent_load(self, pid): + # pylint: disable=unidiomatic-typecheck + if type(pid) is tuple and len(pid) == 3: + if pid[0] == 'scalar': + value = pid[2] + if type(value) not in (str, bytes): + return _unpack_scalar(pid[1], value) + # A string the pickle memoized is copied into a numpy scalar + # once, however many IDs name it, as an array buffer is. + entry = self._strings.get(id(value)) + if entry is None: + entry = (value, _unpack_scalar(pid[1], value)) + self._strings[id(value)] = entry + return entry[1] + if pid[0] == 'npy': + header, payload = pid[1], pid[2] + if type(header) is bytes and type(payload) is bytearray: + # Decoding each buffer once keeps what a row can decode to + # within its own size, and shared arrays shared. Holding + # the pair keeps their ids from being reused. + key = (id(header), id(payload)) + entry = self._arrays.get(key) + if entry is None: + array = _unpack_array(header, payload) + entry = (header, payload, array) + self._arrays[key] = entry + return entry[2] + raise pickle.UnpicklingError('SafeDisk refuses a persistent ID') + + +def unpickler(disk, file): + """Return the unpickler of `disk`, a SafeDisk with ``safe_numpy`` set.""" + return _ArrayUnpickler( + file, + SAFE_PICKLE_NAMES, + disk.safe_types, + disk.safe_modules, + ) + + +def store(disk, numpy, value, key): + """Return the Cache table fields of `value`, which may hold numpy data. + + `disk` is a SafeDisk with ``safe_numpy`` set and `numpy` the imported + module. A value with no numpy in it is stored as an ordinary pickle. + + """ + # pylint: disable=protected-access + try: + disk._check(value, 'value') + except TypeError: + pass + else: + return Disk.store(disk, value, False, key=key) + + buffer = io.BytesIO() + protocol = max(disk.pickle_protocol, _MIN_PROTOCOL) + pickler = _ArrayPickler(buffer, protocol, numpy) + + try: + pickler.dump(value) + except (pickle.PicklingError, AttributeError) as error: + raise TypeError( + 'SafeDisk cannot store value: %s' % (error,) + ) from error + + data = buffer.getvalue() + + # Only a reduced object can look up a name, and scanning the pickle + # would copy every array in it. + if pickler.reduced: + disk._check_pickle(data, value, 'value') + + return disk._store_pickle(data, value, key) diff --git a/diskcache_community/arrays.pyi b/diskcache_community/arrays.pyi new file mode 100644 index 0000000..84a46f0 --- /dev/null +++ b/diskcache_community/arrays.pyi @@ -0,0 +1,8 @@ +from typing import IO, Any + +from diskcache_community.core import SafeDisk, SafeUnpickler + +def unpickler(disk: SafeDisk, file: IO[bytes]) -> SafeUnpickler: ... +def store( + disk: SafeDisk, numpy: Any, value: Any, key: Any +) -> tuple[int, int, str | None, Any]: ... diff --git a/diskcache_community/core.py b/diskcache_community/core.py index 5deec2d..e1b7944 100644 --- a/diskcache_community/core.py +++ b/diskcache_community/core.py @@ -314,9 +314,11 @@ def store(self, value, read, key=UNKNOWN): :return: (size, mode, filename, value) tuple for Cache table """ - # pylint: disable=unidiomatic-typecheck + # pylint: disable=unidiomatic-typecheck,comparison-with-itself type_value = type(value) min_file_size = self.min_file_size + # SQLite stores NaN as NULL, so NaN goes to a pickle below. + raw_float = type_value is float and value == value if ( (type_value is str and len(value) < min_file_size) @@ -324,7 +326,7 @@ def store(self, value, read, key=UNKNOWN): type_value is int and -9223372036854775808 <= value <= 9223372036854775807 ) - or (type_value is float) + or raw_float ): return 0, MODE_RAW, None, value elif type_value is bytes: @@ -347,13 +349,16 @@ def store(self, value, read, key=UNKNOWN): return size, MODE_BINARY, filename, None else: result = pickle.dumps(value, protocol=self.pickle_protocol) + return self._store_pickle(result, value, key) - if len(result) < min_file_size: - return 0, MODE_PICKLE, None, sqlite3.Binary(result) - else: - filename, full_path = self.filename(key, value) - self._write(full_path, result, 'xb') - return len(result), MODE_PICKLE, filename, None + def _store_pickle(self, result, value, key): + """Return Cache table fields for `result`, the pickle of `value`.""" + if len(result) < self.min_file_size: + return 0, MODE_PICKLE, None, sqlite3.Binary(result) + else: + filename, full_path = self.filename(key, value) + self._write(full_path, result, 'xb') + return len(result), MODE_PICKLE, filename, None def _write(self, full_path, iterator_or_bytes, mode, encoding=None): full_dir, _ = op.split(full_path) @@ -605,6 +610,49 @@ def fetch(self, mode, filename, value, read): _PICKLE_PUT_OPS = frozenset({'BINPUT', 'LONG_BINPUT', 'PUT'}) _PICKLE_EXT_OPS = frozenset({'EXT1', 'EXT2', 'EXT4'}) +# Opcodes whose argument is a length-prefixed string or buffer, with the width +# of the length, and the longest one the scan reads. A longer one is an array +# or a large string in a value, which cannot name a module, so it is skipped +# instead of copied. +_PICKLE_BLOB_OPS = { + 'BINBYTES': 4, + 'BINBYTES8': 8, + 'BINUNICODE': 4, + 'BINUNICODE8': 8, + 'BYTEARRAY8': 8, +} +_PICKLE_BLOB_LIMIT = 4096 + + +def _pickle_ops(data): + """Yield (opcode, arg) for the pickle `data`, as pickletools.genops does. + + `data` was written by pickle, so it ends with STOP. The argument of a + blob opcode over _PICKLE_BLOB_LIMIT bytes is None. + + """ + file = io.BytesIO(data) + + while True: + opcode = pickletools.code2op[file.read(1).decode('latin-1')] + width = _PICKLE_BLOB_OPS.get(opcode.name) + + if width is None: + arg = None if opcode.arg is None else opcode.arg.reader(file) + else: + size = int.from_bytes(file.read(width), 'little') + if size > _PICKLE_BLOB_LIMIT: + file.seek(size, io.SEEK_CUR) + arg = None + else: + file.seek(-width, io.SEEK_CUR) + arg = opcode.arg.reader(file) + + yield opcode, arg + + if opcode.name == 'STOP': + return + def _in_modules(module, modules): """Return True when `module` is one of `modules` or inside one.""" @@ -626,7 +674,7 @@ def _pickle_globals(data): pushed = [] memo = {} - for opcode, arg, _ in pickletools.genops(data): + for opcode, arg in _pickle_ops(data): name = opcode.name if name in _PICKLE_EXT_OPS: @@ -800,6 +848,7 @@ class SafeDisk(Disk): safe_types: frozenset = frozenset() safe_modules: tuple = () + safe_numpy: bool = False def __init__(self, directory, min_file_size=0, pickle_protocol=0): super().__init__(directory, min_file_size, pickle_protocol) @@ -820,8 +869,8 @@ def with_types(cls, *types, modules=()): do nothing beyond storing their fields, such as dataclasses, named tuples, enums and pydantic models. Reading looks classes up only in modules that are already imported, so import them before reading. - numpy classes are refused, see ``UNSAFE_MODULES``; store arrays as - bytes, for example with ``numpy.save(..., allow_pickle=False)``. + numpy classes are refused, see ``UNSAFE_MODULES``; use + :meth:`with_numpy` to store arrays. >>> import collections >>> Point = collections.namedtuple('Point', 'x y') @@ -842,7 +891,10 @@ def with_types(cls, *types, modules=()): if not isinstance(entry, type): raise TypeError('with_types takes classes, not %r' % (entry,)) if _in_modules(entry.__module__, UNSAFE_MODULES): - raise ValueError('with_types refuses the class %r' % (entry,)) + raise ValueError( + 'with_types refuses the class %r; use with_numpy() to' + ' store arrays' % (entry,) + ) if isinstance(modules, str): modules = (modules,) @@ -858,7 +910,10 @@ def with_types(cls, *types, modules=()): or _in_modules(unsafe, (prefix,)) for unsafe in UNSAFE_MODULES ): - raise ValueError('with_types refuses the module %r' % (entry,)) + raise ValueError( + 'with_types refuses the module %r; use with_numpy() to' + ' store arrays' % (entry,) + ) prefixes.append(prefix) return _derive_disk( @@ -871,14 +926,48 @@ def with_types(cls, *types, modules=()): }, ) - def _unpickler(self, data): + @classmethod + def with_numpy(cls): + """Return a SafeDisk subclass that also stores numpy arrays. + + :class:`SafeDisk` refuses numpy, because unpickling an array is not + memory-safe (see ``UNSAFE_MODULES``). The returned class writes + arrays through pickle's persistent IDs in numpy's ``.npy`` format + with ``allow_pickle=False``, anywhere in a value, and checks each + header before it reads an array back. numpy scalars keep their type. + Object arrays, ``ndarray`` subclasses and other numpy scalars raise + :exc:`TypeError` when written. It combines with :meth:`with_types` + and :meth:`with_options` in any order. + + A value with no numpy in it is an ordinary pickle. An entry holding + numpy data cannot be read by any other disk, or by diskcache 5.6.3: + they raise :exc:`pickle.UnpicklingError` for it, and a disk with + ``on_error='miss'`` deletes it, so select this disk in every process + that shares the directory. A process without numpy raises + :exc:`ImportError` for such an entry and keeps it. + + >>> ArrayDisk = SafeDisk.with_numpy() + >>> ArrayDisk.safe_numpy + True + + :return: SafeDisk subclass to pass as ``disk=`` + + """ + return _derive_disk(cls, {'safe_numpy': True}) + + def _plain_unpickler(self, file): return SafeUnpickler( - io.BytesIO(data), + file, SAFE_PICKLE_NAMES, self.safe_types, self.safe_modules, ) + def _unpickler(self, file): + if self.safe_numpy: + return arrays.unpickler(self, file) + return self._plain_unpickler(file) + def _extended(self): return bool(self.safe_types or self.safe_modules) @@ -892,12 +981,17 @@ def _check_object(self, obj, what): else: return - data = pickle.dumps(obj, protocol=self.pickle_protocol) + self._check_pickle( + pickle.dumps(obj, protocol=self.pickle_protocol), obj, what + ) + + def _check_pickle(self, data, obj, what): + """Raise TypeError unless the reader would accept the pickle `data`.""" names = _pickle_globals(data) if names is None: try: - self._unpickler(data).load() + self._unpickler(io.BytesIO(data)).load() except pickle.UnpicklingError as error: raise TypeError( 'SafeDisk cannot store %s: %s' % (what, error) @@ -906,7 +1000,7 @@ def _check_object(self, obj, what): for module, name in sorted(names): try: - self._unpickler(b'').find_class(module, name) + self._unpickler(io.BytesIO()).find_class(module, name) except pickle.UnpicklingError: raise TypeError( 'SafeDisk cannot store %s of type %s.%s: it pickles' @@ -1003,12 +1097,17 @@ def put(self, key): def get(self, key, raw): if raw: return super().get(key, raw) - return self._unpickler(key).load() + # Keys never hold numpy, so a key row gets no persistent IDs. + return self._plain_unpickler(io.BytesIO(key)).load() def store(self, value, read, key=UNKNOWN): # pylint: disable=unidiomatic-typecheck - if not read and type(value) not in (bytes, str, int, float): - self._check_object(value, 'value') + if read or type(value) in (bytes, str, int, float): + return super().store(value, read, key=key) + numpy = sys.modules.get('numpy') if self.safe_numpy else None + if numpy is not None: + return arrays.store(self, numpy, value, key) + self._check_object(value, 'value') return super().store(value, read, key=key) def fetch(self, mode, filename, value, read): @@ -1017,8 +1116,9 @@ def fetch(self, mode, filename, value, read): if mode == MODE_PICKLE: if value is None: - value = super().fetch(MODE_BINARY, filename, None, False) - return self._unpickler(value).load() + with super().fetch(MODE_BINARY, filename, None, True) as file: + return self._unpickler(file).load() + return self._unpickler(io.BytesIO(value)).load() return super().fetch(mode, filename, value, read) @@ -1043,10 +1143,14 @@ class SafeJSONDisk(JSONDisk, SafeDisk): The on-disk format is the one :class:`JSONDisk` writes, so a directory written with either opens with the other. Keys and values are limited to what JSON can hold, and with no pickles to read there is no allowlist for - :meth:`SafeDisk.with_types` to extend. + :meth:`SafeDisk.with_types` to extend; :meth:`SafeDisk.with_numpy` raises. """ + @classmethod + def with_numpy(cls): + raise TypeError('SafeJSONDisk stores only JSON, not numpy arrays') + def get(self, key, raw): if not raw: raise pickle.UnpicklingError('SafeJSONDisk refuses a pickled key') @@ -1117,10 +1221,10 @@ def default_disk(): ``DISKCACHE_DISK`` is empty or ``default`` for :class:`Disk`, ``safe`` for :class:`SafeDisk`, or the import path of a Disk subclass, such as - ``myapp.cache:AppDisk``. With ``safe``, ``DISKCACHE_SAFE_TYPES`` lists - classes to allow by import path and ``DISKCACHE_SAFE_MODULES`` lists - module names whose classes to allow, both comma-separated, as - :meth:`SafeDisk.with_types` takes them. + ``myapp.cache:AppDisk``. With ``safe`` or a SafeDisk subclass, + ``DISKCACHE_SAFE_TYPES`` lists classes to allow by import path and + ``DISKCACHE_SAFE_MODULES`` lists module names whose classes to allow, + both comma-separated, as :meth:`SafeDisk.with_types` takes them. :return: Disk type or subclass :raises ValueError: if an environment variable is invalid @@ -1171,27 +1275,26 @@ def _disk_from_environment(name, types, modules): extras = types.strip() or modules.strip() if keyword == 'safe': - classes = [ - _import_object(entry, ENV_SAFE_TYPES) - for entry in _split_list(types) - ] - prefixes = _split_list(modules) - if not (classes or prefixes): - return SafeDisk - return SafeDisk.with_types(*classes, modules=prefixes) - - if extras: + disk = SafeDisk + elif keyword in ('', 'default'): + disk = Disk + else: + disk = _import_object(name, ENV_DISK) + _check_disk_type(disk, ENV_DISK) + + if not extras: + return disk + + if not issubclass(disk, SafeDisk): raise ValueError( - '%s and %s apply only with %s=safe' + '%s and %s apply only with %s=safe or a SafeDisk subclass' % (ENV_SAFE_TYPES, ENV_SAFE_MODULES, ENV_DISK) ) - if keyword in ('', 'default'): - return Disk - - disk = _import_object(name, ENV_DISK) - _check_disk_type(disk, ENV_DISK) - return disk + classes = [ + _import_object(entry, ENV_SAFE_TYPES) for entry in _split_list(types) + ] + return disk.with_types(*classes, modules=_split_list(modules)) def _meta_result(value, expire_time, tag, store_time, meta=(None, None, None)): @@ -1500,9 +1603,19 @@ def after_fork_child(self): # Caches that have logged an unreadable entry as a warning already. _UNREADABLE_LOGGED: 'weakref.WeakSet[Cache]' = weakref.WeakSet() + +class _ReaderMissing(ImportError): + """A disk needs a package that is not installed to read an entry. + + Unlike a pickle that names a class that no longer exists, which raises an + ordinary ImportError, this says nothing about the entry itself. + + """ + + # Errors that on_error='miss' never treats as an unreadable entry: a missing # value file is already a miss, and the others say nothing about the entry. -_NOT_UNREADABLE = (OSError, MemoryError, RecursionError) +_NOT_UNREADABLE = (OSError, MemoryError, RecursionError, _ReaderMissing) # Connections inherited from a parent process, kept open on purpose. _INHERITED_CONNECTIONS: 'list[sqlite3.Connection]' = [] @@ -3922,3 +4035,8 @@ def reset(self, key, value=ENOVAL, update=True): setattr(self, key, value) return value + + +# arrays imports the classes above, so it is imported once they exist. +# pylint: disable=wrong-import-position,cyclic-import +from . import arrays # noqa: E402 # isort: skip diff --git a/diskcache_community/core.pyi b/diskcache_community/core.pyi index a2cea74..2dc0971 100644 --- a/diskcache_community/core.pyi +++ b/diskcache_community/core.pyi @@ -9,6 +9,7 @@ from typing import ( Generator, Generic, Literal, + NoReturn, overload, ) @@ -188,12 +189,17 @@ class SafeUnpickler(pickle.Unpickler): class SafeDisk(Disk): safe_types: frozenset[type] safe_modules: tuple[str, ...] + safe_numpy: bool @classmethod def with_types( cls, *types: type, modules: tuple[str, ...] = ... ) -> type[Self]: ... + @classmethod + def with_numpy(cls) -> type[Self]: ... -class SafeJSONDisk(JSONDisk, SafeDisk): ... +class SafeJSONDisk(JSONDisk, SafeDisk): + @classmethod + def with_numpy(cls) -> NoReturn: ... def configure(disk: type[Disk] | None) -> None: ... def default_disk() -> type[Disk]: ... diff --git a/docs/fix-cve-2025-69872.rst b/docs/fix-cve-2025-69872.rst index fc1a0b4..78bca96 100644 --- a/docs/fix-cve-2025-69872.rst +++ b/docs/fix-cve-2025-69872.rst @@ -111,20 +111,21 @@ Steps With ``DISKCACHE_DISK=safe``, the same lists go into ``DISKCACHE_SAFE_TYPES`` and ``DISKCACHE_SAFE_MODULES``. numpy classes - cannot be allowed, because a crafted numpy pickle can corrupt memory. Store - arrays as bytes instead: + cannot be allowed, because a crafted numpy pickle can corrupt memory. Use + :meth:`SafeDisk.with_numpy ` to + store arrays instead, which saves them in numpy's ``.npy`` format and + combines with ``with_types``: .. code-block:: python - import io - - import numpy + from diskcache_community import Cache, SafeDisk - buffer = io.BytesIO() - numpy.save(buffer, array, allow_pickle=False) - cache['embedding'] = buffer.getvalue() + ArrayDisk = SafeDisk.with_numpy().with_types(Message, Point) + cache = Cache('/var/tmp/mycache', disk=ArrayDisk) + cache['embedding'] = array - array = numpy.load(io.BytesIO(cache['embedding']), allow_pickle=False) + To select it from the environment, name a class like ``ArrayDisk`` in + ``DISKCACHE_DISK``, such as ``myapp.cache:ArrayDisk``. The same lists apply. 5. **Deal with entries already in the cache.** A cache written with the default :class:`Disk ` may hold objects that diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 5afeddc..a461648 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -1013,11 +1013,11 @@ more than store their fields, such as dataclasses, named tuples, enums and pydantic models, because allowing a class lets a crafted pickle create instances of it with any field values. numpy classes are refused: numpy restores a dtype's internal flags from the pickle without checking them, so a -crafted array pickle can corrupt memory. Store arrays as bytes instead, for -example with ``numpy.save(buffer, array, allow_pickle=False)``. A write whose -value pickles a class outside the allowlist raises :exc:`TypeError` that names -the class. Calling ``with_types`` on the returned class adds to its allowlist, -and a cache that uses it can be pickled as usual. +crafted array pickle can corrupt memory. Use :meth:`SafeDisk.with_numpy +` to store arrays instead. A write +whose value pickles a class outside the allowlist raises :exc:`TypeError` that +names the class. Calling ``with_types`` on the returned class adds to its +allowlist, and a cache that uses it can be pickled as usual. .. _tutorial-safejsondisk: @@ -1056,6 +1056,69 @@ extend and :meth:`SafeDisk.with_types ``SafeJSONDisk.with_options(on_error='miss')`` to drop refused entries, as described in :ref:`tutorial-unreadable-entries`. +.. _tutorial-with-numpy: + +Storing numpy arrays +.................... + +:class:`SafeDisk ` refuses numpy, because +unpickling an array is not memory-safe. :meth:`SafeDisk.with_numpy +` returns a subclass that stores +arrays without pickling them, by writing them through pickle's persistent IDs +in numpy's ``.npy`` format: + +.. code-block:: python + + import numpy + + from diskcache_community import Cache, SafeDisk + + ArrayDisk = SafeDisk.with_numpy() + cache = Cache('/tmp/mycache', disk=ArrayDisk) + cache['embedding'] = numpy.arange(3.0) + cache['result'] = {'scores': numpy.zeros(2), 'best': numpy.float64(1.5)} + +* An array anywhere in a value is stored as its ``.npy`` header and its data, + including inside classes allowed with :meth:`SafeDisk.with_types + `, which can be called before or + after ``with_numpy``. The header is checked against the data before it is + read back. An array that appears twice in a value is stored once and is one + array when read back. Writing copies the array once, into the row, as the + default :class:`Disk ` does, and reading takes no + second copy of it. +* A value outside the standard-library types is pickled at protocol 5 or + higher, whatever ``disk_pickle_protocol`` says, so that array data can be + stored in place. Other values use the configured protocol. +* numpy booleans, integers, floats, complex numbers and strings read back as + the same numpy type with the same value, including NaN. +* Arrays of objects, subclasses of ``ndarray`` such as masked arrays and + ``numpy.matrix``, ``datetime64`` and ``timedelta64`` scalars, and arrays that + the reader would refuse, such as a dtype with no width, raise + :exc:`TypeError` when written. + +Keys behave as in :class:`SafeDisk `, so numpy +values in a key raise :exc:`TypeError`, including the arguments of a function +decorated with :meth:`memoize `. + +Everything else behaves as in :class:`SafeDisk +`, and ``with_options(on_error='miss')`` works +the same way. An entry whose array data has been tampered with raises +:exc:`pickle.UnpicklingError` when read, or reads as missing with +``on_error='miss'``. + +A value with no numpy in it is an ordinary pickle. An entry holding numpy data +is not readable by any other disk, including a plain :class:`SafeDisk +`, or by diskcache 5.6.3: they open the +directory and raise :exc:`pickle.UnpicklingError` for that entry, and a disk +with ``on_error='miss'`` deletes it. The choice of disk is not stored in the +cache directory, so select this disk in every process that shares a directory. +To select it from the environment, define it in a module of your own and name +it in ``DISKCACHE_DISK``, such as ``myapp.cache:ArrayDisk``. A process that +cannot import numpy raises :exc:`ImportError` for an entry holding numpy data +and keeps it, even with ``on_error='miss'``. Arrays that the default +:class:`Disk ` pickled earlier are refused like any +other pickle outside the allowlist, so compute them again. + .. _tutorial-default-disk: Default disk @@ -1081,8 +1144,10 @@ or in the environment, without changing any code: DISKCACHE_DISK=safe python app.py ``DISKCACHE_DISK`` takes ``safe``, ``default``, or the import path of a disk -class such as ``myapp.cache:AppDisk``. With ``safe``, ``DISKCACHE_SAFE_TYPES`` -and ``DISKCACHE_SAFE_MODULES`` list classes and module names, comma-separated, +class such as ``myapp.cache:AppDisk``. With ``safe`` or a subclass of +:class:`SafeDisk `, such as one built with +``with_numpy``, ``DISKCACHE_SAFE_TYPES`` and +``DISKCACHE_SAFE_MODULES`` list classes and module names, comma-separated, passed to :meth:`SafeDisk.with_types `. A disk passed explicitly always wins, then :func:`configure `, then the @@ -1112,7 +1177,9 @@ suits a cache whose entries can always be computed again: ``pop``, ``pull``, ``peek`` and ``peekitem`` move on to the next item, and iteration, including the items and values of an :class:`Index `, skips entries that cannot be read. ``len`` and -``in`` still count an unreadable entry until a read deletes it. The first +``in`` still count an unreadable entry until a read deletes it, so code like +``if key in cache: return cache[key]`` can raise :exc:`KeyError` for such an +entry; use ``cache.get(key, default)`` instead. The first unreadable entry of each cache is logged as a warning on the ``diskcache_community`` logger. A read-only cache returns the default from ``get`` without deleting the entry, and its ``peek`` and ``peekitem`` raise, diff --git a/tests/test_core.py b/tests/test_core.py index 5ccba7d..7d2c253 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -5,6 +5,7 @@ import gc import hashlib import io +import math import os import os.path as op import pathlib @@ -2220,3 +2221,27 @@ def test_get_store_time_read(cache): reader, store_time = cache.get('key', read=True, store_time=True) assert reader.read() == b'value' assert isinstance(store_time, float) + + +@pytest.mark.parametrize('disk', [dc.Disk, dc.SafeDisk, dc.JSONDisk]) +def test_nan_value_is_kept(disk): + # SQLite stores NaN as NULL, which would read back as None. + with dc.Cache(disk=disk) as cache: + cache['key'] = float('nan') + cache['inf'] = float('inf') + assert math.isnan(cache['key']) + assert cache['inf'] == float('inf') + assert cache.check() == [] + shutil.rmtree(cache.directory, ignore_errors=True) + + +def test_memoize_nan_is_kept_on_a_hit(): + with dc.Cache() as cache: + + @cache.memoize() + def compute(): + return float('nan') + + assert math.isnan(compute()) + assert math.isnan(compute()) + shutil.rmtree(cache.directory, ignore_errors=True) diff --git a/tests/test_default_disk.py b/tests/test_default_disk.py index 73ebacd..9e9c986 100644 --- a/tests/test_default_disk.py +++ b/tests/test_default_disk.py @@ -1,5 +1,6 @@ """Test diskcache_community.configure and default_disk.""" +import decimal import os.path as op import pickle import shutil @@ -263,3 +264,22 @@ def test_default_disk_not_persisted(directory): with dc.Cache(directory) as cache: assert type(cache.disk) is dc.Disk assert cache['key'] == 'value' + + +def test_safe_lists_apply_to_a_safedisk_subclass(monkeypatch): + monkeypatch.setenv(core.ENV_DISK, 'diskcache_community:SafeJSONDisk') + monkeypatch.setenv(core.ENV_SAFE_TYPES, 'decimal:Context') + monkeypatch.setenv(core.ENV_SAFE_MODULES, 'myapp.schemas') + core._disk_from_environment.cache_clear() + disk = dc.default_disk() + assert issubclass(disk, dc.SafeJSONDisk) + assert decimal.Context in disk.safe_types + assert disk.safe_modules == ('myapp.schemas',) + + +def test_safe_lists_refuse_a_disk_outside_safedisk(monkeypatch): + monkeypatch.setenv(core.ENV_DISK, 'diskcache_community:JSONDisk') + monkeypatch.setenv(core.ENV_SAFE_MODULES, 'myapp.schemas') + core._disk_from_environment.cache_clear() + with pytest.raises(ValueError, match='SafeDisk subclass'): + dc.default_disk() diff --git a/tests/test_safedisk.py b/tests/test_safedisk.py index c6153e6..ac12df1 100644 --- a/tests/test_safedisk.py +++ b/tests/test_safedisk.py @@ -640,3 +640,18 @@ def test_django_disk_option(directory): with pytest.raises(ImportError): dc.DjangoCache(directory, {'DISK': 'diskcache_community.NoSuchDisk'}) + + +def test_file_backed_pickle_reads_from_its_file(directory): + with dc.Cache(directory, disk=dc.SafeDisk, disk_min_file_size=0) as cache: + cache['key'] = {'a': [1, 2]} + assert cache['key'] == {'a': [1, 2]} + assert cache.check() == [] + + +def test_missing_file_for_a_pickle_is_a_miss(directory): + with dc.Cache(directory, disk=dc.SafeDisk, disk_min_file_size=0) as cache: + cache['key'] = {'a': [1, 2]} + filename = cache._sql('SELECT filename FROM Cache').fetchone()[0] + os.remove(op.join(directory, filename)) + assert cache.get('key', 'default') == 'default' diff --git a/tests/test_safedisk_numpy.py b/tests/test_safedisk_numpy.py new file mode 100644 index 0000000..5084d4e --- /dev/null +++ b/tests/test_safedisk_numpy.py @@ -0,0 +1,1403 @@ +"""Test SafeDisk.with_numpy().""" + +import collections as co +import dataclasses +import io +import os.path as op +import pickle +import pickletools +import shutil +import sqlite3 +import struct +import subprocess +import sys +import tempfile +import tracemalloc +import warnings + +import pytest + +import diskcache_community as dc +from diskcache_community import arrays, core +from diskcache_community.core import MODE_PICKLE + +np = pytest.importorskip('numpy') + +pytestmark = pytest.mark.filterwarnings('ignore', category=dc.EmptyDirWarning) + +ArrayDisk = dc.SafeDisk.with_numpy() +MissDisk = ArrayDisk.with_options(on_error='miss') + + +def _execute(path): + """Marker that a crafted pickle ran code: it creates `path`.""" + with open(path, 'w', encoding='utf-8') as writer: + writer.write('executed') + return 'executed' + + +class Poison: + """Object whose pickle calls `_execute` when loaded.""" + + def __init__(self, path): + self.path = path + + def __reduce__(self): + return (_execute, (self.path,)) + + +class Outside: + """A class outside the allowlist.""" + + +class Temporaries: + """Object that hands pickle arrays nothing else refers to.""" + + def __init__(self, count): + self.count = count + + def __reduce__(self): + arrays = (np.arange(n, n + 2.0) for n in range(self.count)) + return (list, (), None, arrays) + + +class State: + """A plain class holding arrays, as a library's state object does.""" + + def __init__(self, tokens, scores): + self.tokens = tokens + self.scores = scores + + +@dataclasses.dataclass +class Result: + name: str + vector: object + + +Pair = co.namedtuple('Pair', 'left right') + + +@pytest.fixture +def directory(): + directory = tempfile.mkdtemp(prefix='diskcache-') + yield directory + shutil.rmtree(directory, ignore_errors=True) + + +def same(actual, expected): + assert type(actual) is np.ndarray + assert actual.dtype == expected.dtype + assert actual.shape == expected.shape + assert actual.tobytes() == expected.tobytes() + + +def npy_parts(array): + """Return the ``.npy`` header of `array` and its data, as numpy saves it.""" + buffer = io.BytesIO() + np.save(buffer, array, allow_pickle=False) + data = buffer.getvalue() + size = len(data) - array.nbytes + return data[:size], data[size:] + + +class Pid: + """Marks where a forged pickle holds the persistent ID `pid`.""" + + def __init__(self, pid): + self.pid = pid + + +class Forger(pickle.Pickler): + """Pickler for rows that does not use the code tested.""" + + def persistent_id(self, obj): + return obj.pid if isinstance(obj, Pid) else None + + +def forge(value): + buffer = io.BytesIO() + Forger(buffer, protocol=5).dump(value) + return buffer.getvalue() + + +def persid(pid): + return forge(Pid(pid)) + + +def array_pid(parts): + header, data = parts + return ('npy', header, bytearray(data)) + + +def write_row(directory, data): + """Store the pickle `data` in the value column of a row.""" + with dc.Cache(directory, disk=dc.SafeDisk) as cache: + cache['key'] = 'value' + con = sqlite3.connect(op.join(directory, 'cache.db')) + con.execute( + 'UPDATE Cache SET mode = ?, value = ?', + (MODE_PICKLE, sqlite3.Binary(data)), + ) + con.commit() + con.close() + + +def row_bytes(cache, key): + return bytes( + cache._sql('SELECT value FROM Cache WHERE key = ?', (key,)).fetchone()[ + 0 + ] + ) + + +def npy(descr, shape, body=b'', version=1): + header = "{'descr': %s, 'fortran_order': False, 'shape': %s, }" % ( + descr, + shape, + ) + return npy_header(header, body, version) + + +def npy_header(header, body=b'', version=1): + header = header.encode('latin1') + size = ' pickle.HIGHEST_PROTOCOL: + pytest.skip('protocol not available') + + with dc.Cache( + directory, disk=ArrayDisk, disk_pickle_protocol=protocol + ) as cache: + cache['key'] = {'a': np.arange(2.0), 'b': np.int64(1)} + same(cache['key']['a'], np.arange(2.0)) + assert cache['key']['b'] == 1 + + +@pytest.mark.parametrize( + 'scalar', + [ + np.float64(1.5), + np.float32(0.1), + np.float16(1.5), + np.int8(-2), + np.int16(2), + np.int32(2), + np.int64(2), + np.uint8(2), + np.uint16(2), + np.uint32(2), + np.uint64(2**64 - 1), + np.bool_(True), + np.complex64(1j), + np.complex128(1j), + np.str_('text'), + np.bytes_(b'data'), + ], + ids=lambda value: type(value).__name__, +) +def test_scalars_keep_their_type_and_value(directory, scalar): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = scalar + cache['nested'] = {'metrics': [scalar, scalar]} + assert cache['key'] == scalar + assert type(cache['key']) is type(scalar) + assert type(cache['nested']['metrics'][1]) is type(scalar) + + +def test_scalars_of_one_value_keep_their_own_types(directory): + """Python caches small ints, so equal scalars must not share a type.""" + scalars = [ + np.int8(2), + np.int64(2), + np.uint8(2), + np.float64(2.0), + np.float32(2.0), + np.bool_(True), + np.int64(1), + np.str_(''), + np.bytes_(b''), + ] + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = scalars + result = cache['key'] + + assert [type(item) for item in result] == [type(item) for item in scalars] + assert result == scalars + + +@pytest.mark.parametrize('kind', [np.float64, np.float32, np.float16]) +def test_scalar_nan_and_infinity_are_kept(directory, kind): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['nan'] = kind('nan') + cache['inf'] = kind('inf') + cache['nested'] = [kind('nan')] + assert type(cache['nan']) is kind + assert np.isnan(cache['nan']) + assert cache['inf'] == kind('inf') + assert np.isnan(cache['nested'][0]) + + +def test_scalars_as_dict_keys_keep_their_type(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = {np.int64(1): 'one'} + assert cache['key'] == {1: 'one'} + assert type(next(iter(cache['key']))) is np.int64 + + +@pytest.mark.parametrize( + 'make', + [ + lambda scalar: {scalar}, + lambda scalar: frozenset([scalar]), + lambda scalar: co.deque([scalar]), + lambda scalar: co.OrderedDict(a=scalar), + lambda scalar: co.Counter({'a': scalar}), + lambda scalar: co.defaultdict(int, a=scalar), + ], + ids=['set', 'frozenset', 'deque', 'ordereddict', 'counter', 'defaultdict'], +) +def test_scalars_in_the_other_safe_containers(directory, make): + value = make(np.int64(2)) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = value + result = cache['key'] + + assert type(result) is type(value) + item = next(iter(result.values() if hasattr(result, 'values') else result)) + assert item == 2 + assert type(item) is np.int64 + + +@pytest.mark.parametrize( + 'scalar', + [ + np.datetime64('2026-01-01T00:00:00.000000000'), + np.timedelta64(5, 'ns'), + ], + ids=lambda value: type(value).__name__, +) +def test_time_scalars_are_refused(directory, scalar): + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(TypeError): + cache['key'] = scalar + with pytest.raises(TypeError): + cache['key'] = [scalar] + + +def test_long_double_reads_back_or_is_refused(directory): + value = np.longdouble(1.5) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + try: + cache['key'] = {'a': value} + except TypeError: + return + assert cache['key'] == {'a': 1.5} + + +@pytest.mark.parametrize( + 'array', + [ + np.array([object()], dtype=object), + np.array([(1, None)], dtype=[('a', 'i4'), ('b', 'O')]), + np.zeros(1, dtype=[('a', 'O', (2,))]), + ], + ids=['object', 'field', 'subarray'], +) +def test_object_arrays_are_refused(directory, array): + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(TypeError): + cache['key'] = array + assert len(cache) == 0 + + +@pytest.mark.parametrize( + 'array', + [ + np.zeros(1, dtype=[('α', 'i4')]), + np.zeros(1, dtype=[('f%d' % number, 'i4') for number in range(900)]), + np.zeros(3, dtype=[]), + np.zeros( + 2, + dtype={ + 'names': ['a', 'b'], + 'formats': ['i4', 'i2'], + 'offsets': [0, 2], + }, + ), + ], + ids=['format-3', 'long-header', 'zero-width', 'overlapping-fields'], +) +def test_arrays_the_reader_would_refuse_fail_at_write(directory, array): + with dc.Cache(directory, disk=MissDisk) as cache: + with pytest.raises(TypeError): + cache['key'] = array + assert len(cache) == 0 + + +@pytest.mark.parametrize( + 'make', + [ + lambda: np.matrix([1, 2]), + lambda: np.ma.masked_array([1, 2], mask=[1, 0]), + lambda: np.array([1.0]).view(np.recarray), + ], + ids=['matrix', 'masked', 'recarray'], +) +def test_subclasses_are_refused(directory, make): + """numpy.save would silently drop what makes them a subclass.""" + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(TypeError, match='subclass'): + cache['key'] = make() + with pytest.raises(TypeError, match='subclass'): + cache['key'] = [make()] + + +def test_numpy_in_a_key_is_refused(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(TypeError): + cache[(np.int64(1),)] = 'value' + with pytest.raises(TypeError): + cache[(np.arange(2.0).tobytes(), np.arange(2.0))] = 'value' + + +def test_memoize(directory): + calls = [] + + with dc.Cache(directory, disk=ArrayDisk) as cache: + + @cache.memoize() + def compute(size): + calls.append(size) + return np.arange(float(size)) + + same(compute(3), np.arange(3.0)) + same(compute(3), np.arange(3.0)) + + assert calls == [3] + + +def test_memoize_scalar_keeps_its_type_on_a_hit(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + + @cache.memoize() + def compute(): + return np.array([np.nan, 1.0]).mean() + + first, second = compute(), compute() + + assert type(first) is type(second) is np.float64 + assert np.isnan(first) and np.isnan(second) + + +def test_memoize_scalar_argument_is_refused(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + + @cache.memoize() + def compute(value): + return value + + with pytest.raises(TypeError): + compute(np.int64(3)) + + +def test_plain_values_behave_as_in_safedisk(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + value = {'a': [1, 2.5, 'x', None, (True,)], 'b': {'c': b'data'}} + cache['key'] = value + assert cache['key'] == value + with pytest.raises(TypeError): + cache['outside'] = Outside() + with pytest.raises(TypeError): + cache['function'] = lambda: None + + +def test_plain_values_are_ordinary_pickles(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = {'a': [1, 2]} + names = { + op.name for op, _, _ in pickletools.genops(row_bytes(cache, 'key')) + } + + assert 'BINPERSID' not in names + + with dc.Cache(directory, disk=dc.Disk) as cache: + assert cache['key'] == {'a': [1, 2]} + + +def test_plain_values_skip_the_array_pickler(directory, monkeypatch): + """Otherwise every object in them costs a call into Python.""" + + def refuse(*args): + raise AssertionError('plain value reached the array pickler') + + monkeypatch.setattr(arrays, '_ArrayPickler', refuse) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = {'a': [1, 2.5, 'x']} + assert cache['key'] == {'a': [1, 2.5, 'x']} + + +def test_plain_values_do_not_import_numpy(directory): + code = ( + 'import sys\n' + 'import diskcache_community as dc\n' + 'assert "numpy" not in sys.modules\n' + 'disk = dc.SafeDisk.with_numpy()\n' + 'with dc.Cache(sys.argv[1], disk=disk) as cache:\n' + ' cache["plain"] = {"a": [1, 2]}\n' + ' assert cache["plain"] == {"a": [1, 2]}\n' + 'assert "numpy" not in sys.modules\n' + ) + subprocess.run([sys.executable, '-c', code, directory], check=True) + + +def test_reads_an_array_in_a_process_that_has_not_imported_numpy(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = np.arange(2.0) + + code = ( + 'import sys\n' + 'import diskcache_community as dc\n' + 'assert "numpy" not in sys.modules\n' + 'disk = dc.SafeDisk.with_numpy()\n' + 'with dc.Cache(sys.argv[1], disk=disk) as cache:\n' + ' assert cache["key"].tolist() == [0.0, 1.0]\n' + ) + subprocess.run([sys.executable, '-c', code, directory], check=True) + + +def test_without_numpy_an_entry_raises_and_is_kept(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = np.arange(2.0) + + code = ( + 'import sys\n' + 'sys.modules["numpy"] = None\n' + 'import diskcache_community as dc\n' + 'disk = dc.SafeDisk.with_numpy().with_options(on_error="miss")\n' + 'with dc.Cache(sys.argv[1], disk=disk) as cache:\n' + ' try:\n' + ' cache.get("key")\n' + ' except ImportError as error:\n' + ' assert "numpy" in str(error)\n' + ' else:\n' + ' raise AssertionError("no ImportError")\n' + ' cache.expire()\n' + ' assert len(cache) == 1\n' + ) + subprocess.run([sys.executable, '-c', code, directory], check=True) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + same(cache['key'], np.arange(2.0)) + + +def test_missing_numpy_is_not_an_unreadable_entry(directory, monkeypatch): + with dc.Cache(directory, disk=MissDisk) as cache: + cache['key'] = np.arange(2.0) + monkeypatch.setitem(sys.modules, 'numpy', None) + with pytest.raises(ImportError, match='needs numpy'): + cache.get('key') + monkeypatch.undo() + same(cache['key'], np.arange(2.0)) + + +def test_scalar_without_numpy_raises_and_is_kept(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = np.float64(1.5) + + code = ( + 'import sys\n' + 'sys.modules["numpy"] = None\n' + 'import diskcache_community as dc\n' + 'disk = dc.SafeDisk.with_numpy().with_options(on_error="miss")\n' + 'with dc.Cache(sys.argv[1], disk=disk) as cache:\n' + ' try:\n' + ' cache.get("key")\n' + ' except ImportError:\n' + ' pass\n' + ' else:\n' + ' raise AssertionError("no ImportError")\n' + ' assert len(cache) == 1\n' + ) + subprocess.run([sys.executable, '-c', code, directory], check=True) + + +def test_arrays_read_back_writeable(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = np.arange(3.0) + result = cache['key'] + + assert result.flags.writeable + result[0] = 9.0 + assert result.tolist() == [9.0, 1.0, 2.0] + + +def test_fortran_order_round_trip(directory): + array = np.arange(6.0).reshape(2, 3, order='F') + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = array + assert cache['key'].tolist() == array.tolist() + + +def test_empty_array_of_a_dtype_with_no_width(directory): + array = np.zeros(0, dtype=[]) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = array + assert cache['key'].shape == (0,) + + +def test_writing_a_large_array_copies_it_once(directory): + """pickle copies the array's own buffer into the row, and nothing else + copies it, as on the default Disk.""" + array = np.arange(4000000.0) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + tracemalloc.start() + cache['key'] = {'array': array} + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert peak < array.nbytes * 1.3 + + +def test_writing_an_array_in_an_allowed_class_copies_it_once(directory): + """The write check reads past the array instead of copying it.""" + array = np.arange(4000000.0) + + with dc.Cache(directory, disk=ArrayDisk.with_types(Result)) as cache: + tracemalloc.start() + cache['key'] = Result('name', array) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + same(cache['key'].vector, array) + + assert peak < array.nbytes * 1.3 + + +def test_reading_a_large_array_takes_no_second_copy(directory): + array = np.arange(4000000.0) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = array + tracemalloc.start() + result = cache['key'] + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert result.sum() == array.sum() + assert peak < array.nbytes * 1.2 + + +def test_read_true_still_returns_a_file(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache.set('key', io.BytesIO(b'raw'), read=True) + with cache.get('key', read=True) as reader: + assert reader.read() == b'raw' + + +def test_large_array_lives_in_a_file(directory): + array = np.arange(100000.0) + + with dc.Cache(directory, disk=ArrayDisk, disk_min_file_size=0) as cache: + cache['key'] = array + same(cache['key'], array) + assert cache.check() == [] + + +def test_fanout_index_and_deque(directory): + array = np.arange(3.0) + + with dc.FanoutCache(op.join(directory, 'fanout'), disk=ArrayDisk) as cache: + cache['key'] = array + same(cache['key'], array) + + index = dc.Index(op.join(directory, 'index'), disk=ArrayDisk) + index['key'] = array + same(index['key'], array) + index.cache.close() + + deque = dc.Deque(directory=op.join(directory, 'deque'), disk=ArrayDisk) + deque.append(array) + same(deque.pop(), array) + deque.cache.close() + + +def test_with_types_and_options_combine(directory): + disk = ArrayDisk.with_types(Pair).with_options(on_error='miss') + + with dc.Cache(directory, disk=disk) as cache: + cache['pair'] = Pair(1, 2) + cache['array'] = np.arange(2.0) + assert cache['pair'] == Pair(1, 2) + same(cache['array'], np.arange(2.0)) + copy = pickle.loads(pickle.dumps(cache)) + assert isinstance(copy.disk, dc.SafeDisk) + assert copy.disk.safe_numpy + assert copy.disk.on_error == 'miss' + copy.close() + + +def test_with_types_still_refuses_numpy(): + with pytest.raises(ValueError, match='with_numpy'): + ArrayDisk.with_types(np.ndarray) + with pytest.raises(ValueError, match='with_numpy'): + dc.SafeDisk.with_types(modules='numpy') + + +def test_with_numpy_combines_in_either_order(directory): + disks = { + 'first': dc.SafeDisk.with_types(Pair).with_numpy(), + 'second': dc.SafeDisk.with_numpy().with_types(Pair), + } + + for name, disk in disks.items(): + assert disk.safe_numpy + assert Pair in disk.safe_types + with dc.Cache(op.join(directory, name), disk=disk) as cache: + cache['key'] = Pair(np.arange(2.0), 'x') + same(cache['key'].left, np.arange(2.0)) + assert cache['key'].right == 'x' + + +def test_with_numpy_twice_adds_nothing(): + assert dc.SafeDisk.safe_numpy is False + assert ArrayDisk.with_numpy()._disk_options == ArrayDisk._disk_options + + +def test_safedisk_still_refuses_numpy(directory): + with dc.Cache(directory, disk=dc.SafeDisk) as cache: + with pytest.raises(TypeError, match='SafeDisk cannot store value'): + cache['key'] = np.arange(2.0) + with pytest.raises(TypeError, match='SafeDisk cannot store value'): + cache['key'] = {'a': np.float64(1.5)} + assert len(cache) == 0 + + +def test_safejsondisk_refuses_with_numpy(): + with pytest.raises(TypeError, match='SafeJSONDisk'): + dc.SafeJSONDisk.with_numpy() + + +def test_environment_selects_a_with_numpy_disk(monkeypatch): + monkeypatch.setenv(core.ENV_DISK, 'tests.test_safedisk_numpy:ArrayDisk') + monkeypatch.setenv(core.ENV_SAFE_TYPES, 'tests.test_safedisk_numpy:Pair') + core._disk_from_environment.cache_clear() + try: + disk = dc.default_disk() + finally: + core._disk_from_environment.cache_clear() + assert issubclass(disk, dc.SafeDisk) + assert disk.safe_numpy + assert Pair in disk.safe_types + + +def test_row_format_is_pinned(directory): + """Changing what is on disk breaks every cache: this is hand-built.""" + array = np.arange(3.0) + header = ( + b"\x93NUMPY\x01\x00v\x00{'descr': '