From 17b827eb953e6afe4ee6fec2c03fb56472f50944 Mon Sep 17 00:00:00 2001 From: August Vi Date: Thu, 5 Nov 2026 18:41:29 -0800 Subject: [PATCH 1/2] Make SafeDisk.with_numpy() array reads as fast as pickle, and store numpy arrays, scalars, time scalars, object, record and masked arrays and memory maps in its keys and values. --- CHANGELOG.md | 15 +- CONTRIBUTING.md | 2 +- NOTICE | 9 - README.md | 21 +- diskcache_community/__init__.py | 4 +- diskcache_community/arrays.py | 548 ++++++++++++++++++---- diskcache_community/arrays.pyi | 4 + diskcache_community/core.py | 60 ++- docs/conf.py | 6 +- docs/index.rst | 9 +- docs/tutorial.rst | 47 +- pyproject.toml | 9 +- tests/benchmark_core.py | 7 - tests/benchmark_djangocache.py | 6 - tests/plot.py | 1 - tests/test_safedisk_numpy.py | 788 ++++++++++++++++++++++++++++++-- 16 files changed, 1334 insertions(+), 202 deletions(-) delete mode 100644 NOTICE diff --git a/CHANGELOG.md b/CHANGELOG.md index d7931fa..0054740 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,19 @@ diskcache-community forked from DiskCache at upstream commit ebfa37c (3 March 2024), two commits after the 5.6.3 release. Each entry names the pull request in this repository. Imported upstream pull requests are referenced by number, and their commits keep their original authors in the git history. +## Unreleased + +### Changes + +- `SafeDisk.with_numpy()` keeps the `.npy` headers it has parsed or written, so reading an array no longer parses its header each time. A hit on a 384-value float32 array takes about as long as a pickle with the default `Disk`, where it took about three times as long, and a write is about 30% faster. Each header is still checked against the size of the data, and the stored format is unchanged, so 5.8.2 and 5.8.3 read the new entries and the other way round. +- `SafeDisk.with_numpy()` accepts numpy booleans, integers, floats, complex numbers, strings and bytes in keys, also inside tuples, lists, dicts, sets and frozensets, and stores them as the Python values `item()` gives, so `int64(1)` and `1` name the same entry. They raised `TypeError`. +- `SafeDisk.with_numpy()` accepts arrays in keys, alone or inside tuples and lists, and as `memoize` arguments. Every array in a key is stored in C order, so an equal array of the same dtype and shape names the same entry. A key without numpy keeps the bytes earlier versions wrote. +- `SafeDisk.with_numpy()` stores `datetime64` and `timedelta64` scalars with their unit, arrays of objects element by element (each element checked like any value), record arrays and masked arrays with their type, and memory maps as the plain array of their data. They raised `TypeError`. These entries, and keys holding numpy data, read as unreadable on 5.8.3 and earlier. + +### Fixes + +- `SafeDisk.with_numpy()` stores each short-lived array a value hands to pickle with its own data. A value whose pickling made temporary views of one array, such as an object whose `__reduce__` yields slices, could be stored with one view's data in place of later ones, which then read back wrong. Affected 5.8.2 and 5.8.3. + ## 5.8.3 (2026-11-03) Fixes a regression from 5.8.2 and several `SafeDisk` issues found by testing downstream libraries, and updates the documentation. The `cache.db` schema and the file layout are unchanged. @@ -86,7 +99,7 @@ The package is now imported as `diskcache_community`, so it installs next to the ### Packaging - **Breaking:** the package is imported as `diskcache_community` instead of `diskcache`, so it installs next to the upstream `diskcache` distribution without overwriting it, and a project can move to the fork while other dependencies still require `diskcache`. The public API, the `cache.db` schema, the file layout and the serialization are unchanged, so cache directories are shared with 5.6.3 in both directions. Django settings name `diskcache_community.DjangoCache` and `diskcache_community.SafeDisk`. The import-time warning about an upstream `diskcache` installed in the same environment is gone, and a CI job installs both distributions together, runs the golden tests in both directions and checks that uninstalling either leaves the other working (PR #22). -- `LICENSE` now holds the full Apache License 2.0 text instead of the short boilerplate notice, so GitHub and package indexes detect the license; the copyright and fork notice stay in `NOTICE` (PR #21). +- `LICENSE` now holds the full Apache License 2.0 text instead of the short boilerplate notice, so GitHub and package indexes detect the license. ## 5.7.0 (2026-10-11) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index fe131ff..9d024b4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,6 +1,6 @@ # Contributing -diskcache-community is a community fork of [DiskCache](https://github.com/grantjenks/python-diskcache). Bug reports, fixes and small, focused features are welcome. For anything larger, open an issue first so the approach can be agreed before the work is done. +diskcache-community is a fork of [DiskCache](https://github.com/grantjenks/python-diskcache). Bug reports, fixes and small, focused features are welcome. For anything larger, open an issue first so the approach can be agreed before the work is done. ## Setting up diff --git a/NOTICE b/NOTICE deleted file mode 100644 index e1da452..0000000 --- a/NOTICE +++ /dev/null @@ -1,9 +0,0 @@ -diskcache-community -Copyright 2026 diskcache-community contributors - -This project is a fork of DiskCache (https://github.com/grantjenks/python-diskcache). -Original work copyright 2016-2023 Grant Jenks, licensed under the Apache License, -Version 2.0 (see LICENSE). - -Modifications made in this fork are copyright 2026 diskcache-community contributors -and are distributed under the same license. diff --git a/README.md b/README.md index e7b34ee..1efb5e8 100644 --- a/README.md +++ b/README.md @@ -27,16 +27,6 @@


-> [!IMPORTANT] -> This is a community fork of [python-diskcache](https://github.com/grantjenks/python-diskcache), -> whose last release was 5.6.3 in August 2023. It is published on PyPI as -> `diskcache-community` and imported as `diskcache_community`, so it installs -> next to the upstream `diskcache` package. It carries the fix path for -> [CVE-2025-69872](https://github.com/advisories/GHSA-w8v5-vhqr-4h9v) and -> reads and writes the same cache directories as 5.6.3. -> See [Why this fork](#why-this-fork), [Migrating from diskcache](#migrating-from-diskcache) -> and [Fix CVE-2025-69872 in diskcache](https://diskcache-community.readthedocs.io/en/latest/fix-cve-2025-69872.html). - ## Install ```sh @@ -127,8 +117,7 @@ The upstream project stopped releasing after 5.6.3 (August 2023) and its last commit landed in March 2024. In 2025, CVE-2025-69872 was published against every released version: the cache reads pickled values with `pickle.load`, so anyone who can write into the cache directory can run code in the process -that reads from it. The package still works, but scanners flag it and there -is no fixed upstream version. +that reads from it. This fork publishes the same library as `diskcache-community`. It keeps the `cache.db` schema, the file layout, the serialization modes and the memoize @@ -141,9 +130,6 @@ dependencies in the same environment still use the upstream `diskcache`. The [changelog](CHANGELOG.md) lists every change and [SECURITY.md](SECURITY.md) describes what is and is not fixed. -DiskCache was written and maintained by Grant Jenks from 2016 to 2023. This -fork keeps his license, copyright and author credit; see [NOTICE](NOTICE). - ## Migrating from diskcache The public API and the files on disk are unchanged; only the import name @@ -222,6 +208,5 @@ how to run the tests. Report vulnerabilities privately as described in ## License -Copyright 2016-2023 Grant Jenks. Modifications in this fork are copyright -2026 diskcache-community contributors. Licensed under the -[Apache License, Version 2.0](LICENSE). +Copyright 2026 August V. +Licensed under the [Apache License, Version 2.0](LICENSE). diff --git a/diskcache_community/__init__.py b/diskcache_community/__init__.py index f132cf1..c49978a 100644 --- a/diskcache_community/__init__.py +++ b/diskcache_community/__init__.py @@ -74,6 +74,6 @@ __title__ = 'diskcache-community' __version__ = '5.8.3' __build__ = 0x050803 -__author__ = 'Grant Jenks' +__author__ = 'August V.' __license__ = 'Apache 2.0' -__copyright__ = 'Copyright 2016-2023 Grant Jenks' +__copyright__ = 'Copyright 2026 August V.' diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py index a1348de..a58af6c 100644 --- a/diskcache_community/arrays.py +++ b/diskcache_community/arrays.py @@ -3,9 +3,20 @@ import io import math import pickle +import pickletools +import re +import sqlite3 +import sys import warnings -from .core import SAFE_PICKLE_NAMES, Disk, SafeUnpickler, _ReaderMissing +from .core import ( + SAFE_PICKLE_NAMES, + Disk, + SafeUnpickler, + _in_modules, + _pickle_globals, + _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 @@ -34,6 +45,14 @@ 'bytes': ('bytes_', bytes), } +# The names of datetime64 and timedelta64 scalars, as dtype.name gives them: +# the type, then an optional unit with an optional multiple. The value stored +# is the scalar's int64 count of units, NaT included. +_TIME_NAME = re.compile( + r'(?:datetime64|timedelta64)' + r'(?:\[(?:[1-9][0-9]{0,8})?(?:Y|M|W|D|h|m|s|ms|us|ns|ps|fs|as)\])?\Z' +) + def _import_numpy(): # pylint: disable=import-error,import-outside-toplevel @@ -46,12 +65,31 @@ def _import_numpy(): return numpy -def _read_header(numpy, header, size): +# Headers already parsed, by their bytes, and headers already written, by +# dtype, shape and order. The arrays in a cache mostly share a few of each, +# such as one embedding size, and numpy's header parser, which evaluates the +# header as a Python literal, took most of the time of reading an array. Only +# headers that parse and pass the format checks are kept, only short ones, +# and only for dtypes without fields, since a caller can rename the fields of +# the dtype an array was read with. The size of the data is checked against +# the header on every read. Each table is emptied when full. +_HEADERS_KEPT = 256 +_HEADER_SIZE_KEPT = 256 +_parsed_headers = {} +_written_headers = {} + + +def _keep(table, key, value): + if len(table) >= _HEADERS_KEPT: + table.clear() + table[key] = value + + +def _parse_header(numpy, header): """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 + Raise UnpicklingError unless the header holds one array of a dtype + without objects. 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. @@ -79,11 +117,31 @@ def _read_header(numpy, header, size): 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, count * dtype.itemsize + + +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. + + """ + fields = _parsed_headers.get(header) + if fields is None: + fields = _parse_header(numpy, header) + if len(header) <= _HEADER_SIZE_KEPT and fields[2].names is None: + _keep(_parsed_headers, header, fields) + + shape, fortran, dtype, nbytes = fields + if nbytes != size: + raise pickle.UnpicklingError('SafeDisk refuses the array header') + return shape, fortran, dtype @@ -103,6 +161,21 @@ def _pack_array(numpy, array): if not (flags.writeable and (flags.c_contiguous or flags.f_contiguous)): array = array.copy(order='A') + dtype = array.dtype + key = (dtype, array.shape, not array.flags.c_contiguous) + header = _written_headers.get(key) + if header is None: + header = _write_header(numpy, array) + if len(header) <= _HEADER_SIZE_KEPT and dtype.names is None: + _keep(_written_headers, key, header) + + data = array.reshape(-1, order='A').view(numpy.uint8) + return header, pickle.PickleBuffer(data) + + +def _write_header(numpy, array): + """Return the ``.npy`` header of `array`, or raise TypeError.""" + # pylint: disable=no-member buffer = io.BytesIO() try: with warnings.catch_warnings(): @@ -119,13 +192,13 @@ def _pack_array(numpy, array): message = 'SafeDisk cannot store an array of dtype %s' % (array.dtype,) if array.dtype.hasobject: message += ( - '; a pandas DataFrame holds its labels and text columns in' - ' such arrays, so store its records or columns instead' + '; arrays of objects are stored, but not arrays with object' + ' fields, such as DataFrame.to_records() of a frame with text' + ' columns, so store its columns instead' ) raise TypeError(message) from error - data = array.reshape(-1, order='A').view(numpy.uint8) - return header, pickle.PickleBuffer(data) + return header def _unpack_array(header, data): @@ -153,6 +226,14 @@ def _unpack_array(header, data): def _scalar_name(scalar): """Return the name to store numpy `scalar` under, or raise TypeError.""" dtype = scalar.dtype + if dtype.kind in 'Mm': + # The reader accepts only the units _TIME_NAME matches. + if not _TIME_NAME.match(dtype.name): + raise TypeError( + 'SafeDisk cannot store a scalar of type numpy.%s with the' + ' unit of %s' % (type(scalar).__name__, dtype.name) + ) + return dtype.name name = {'U': 'str', 'S': 'bytes'}.get(dtype.kind, dtype.name) if name not in _SCALARS: raise TypeError( @@ -162,15 +243,43 @@ def _scalar_name(scalar): return name +def _scalar_value(numpy, scalar): + """Return the Python value to store numpy `scalar` as.""" + if scalar.dtype.kind in 'Mm': + return int(scalar.view(numpy.int64)) + value = scalar.item() + if isinstance(value, numpy.generic): + raise TypeError( + 'SafeDisk cannot store a scalar of type numpy.%s' + % (type(scalar).__name__,) + ) + return value + + 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. + `n` zero bytes. A datetime64 or timedelta64 name must match _TIME_NAME + and its value be an int64 count. """ # pylint: disable=broad-except,unidiomatic-typecheck + if type(name) is str and _TIME_NAME.match(name): + if type(value) is not int or not -(2**63) <= value < 2**63: + raise pickle.UnpicklingError('SafeDisk refuses a scalar') + numpy = _import_numpy() + # Every name _TIME_NAME matches builds; the guard is for numpy + # versions that refuse some unit. + try: + count = numpy.array(value, dtype=numpy.int64) + return count.view(numpy.dtype(name))[()] + except Exception as error: # pragma: no cover + raise pickle.UnpicklingError( + 'SafeDisk refuses a scalar' + ) from error + 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') @@ -183,15 +292,132 @@ def _unpack_scalar(name, value): raise pickle.UnpicklingError('SafeDisk refuses a scalar') from error +def _shape(shape): + """Return `shape` if it is a tuple of at most 64 sizes, else None.""" + # pylint: disable=unidiomatic-typecheck + if ( + type(shape) is tuple + and len(shape) <= 64 + and all(type(size) is int and size >= 0 for size in shape) + ): + return shape + return None + + +# The persistent IDs that a row can hold, by kind: +# +# ('npy', header, data) a plain array +# ('scalar', name, value) a numpy scalar +# ('objects', shape, items) an array of objects, in C order +# ('records', array) a numpy.recarray of `array` +# ('masked', data, mask, fill_value, hard) a numpy.ma.MaskedArray +# +# The fields of the last three are pickled as any value is, so the arrays in +# them are persistent IDs of their own and the items are checked like any +# value. The functions below build those three from fields a row chose. + + +def _object_array(shape, items): + """Return the array of objects of `shape` holding `items` in C order.""" + # pylint: disable=broad-except,unidiomatic-typecheck + if _shape(shape) is None or type(items) is not list: + raise pickle.UnpicklingError('SafeDisk refuses an array of objects') + if math.prod(shape) != len(items): + raise pickle.UnpicklingError('SafeDisk refuses an array of objects') + + numpy = _import_numpy() + + try: + array = numpy.empty(shape, dtype=object) + flat = array.reshape(-1) + for index, item in enumerate(items): + flat[index] = item + except Exception as error: + raise pickle.UnpicklingError( + 'SafeDisk refuses an array of objects' + ) from error + return array + + +def _record_array(array): + """Return `array`, a plain array, as a numpy.recarray.""" + numpy = _import_numpy() + # pylint: disable=unidiomatic-typecheck + if type(array) is not numpy.ndarray: + raise pickle.UnpicklingError('SafeDisk refuses a record array') + return array.view(numpy.recarray) + + +def _masked_array(data, mask, fill_value, hard_mask): + """Return the masked array of `data`, `mask` and `fill_value`. + + `mask` is None for no mask and `fill_value` None for numpy's default, + otherwise a 0-d array that numpy checks against the data. + + """ + # pylint: disable=broad-except,unidiomatic-typecheck + numpy = _import_numpy() + ma = numpy.ma + + if ( + type(data) not in (numpy.ndarray, numpy.recarray) + or not ( + fill_value is None + or (type(fill_value) is numpy.ndarray and fill_value.shape == ()) + ) + or type(hard_mask) is not bool + ): + raise pickle.UnpicklingError('SafeDisk refuses a masked array') + + if mask is None: + mask = ma.nomask + elif ( + type(mask) is not numpy.ndarray + or mask.shape != data.shape + or mask.dtype != ma.make_mask_descr(data.dtype) + ): + raise pickle.UnpicklingError('SafeDisk refuses a masked array') + + try: + return ma.MaskedArray( + data, + mask=mask, + fill_value=fill_value, + hard_mask=hard_mask, + copy=False, + ) + except Exception as error: + raise pickle.UnpicklingError( + 'SafeDisk refuses a masked array' + ) from error + + +# The kinds built from fields, with the number of items in their IDs. +_BUILT_KINDS = { + 'objects': (_object_array, 3), + 'records': (_record_array, 2), + 'masked': (_masked_array, 5), +} + + class _ArrayPickler(pickle.Pickler): - """Pickler that writes arrays and numpy scalars as persistent IDs.""" + """Pickler that writes numpy arrays and scalars as persistent IDs. - def __init__(self, file, protocol, numpy): + With `c_order` set, as for keys, every array is written in C order, so + arrays equal in dtype, shape and data give equal bytes. + + """ + + def __init__(self, file, protocol, numpy, c_order=False): super().__init__(file, protocol=protocol) self._numpy = numpy - self._arrays = {} + self._ma = sys.modules.get('numpy.ma') + self._c_order = c_order + self._ids = {} self._scalar_names = {} self.reduced = False + self.wrote_numpy = False + self.wrote_objects = False def reducer_override(self, obj): """Note that an object was reduced, which can look up a name. @@ -207,37 +433,74 @@ def reducer_override(self, obj): 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 + # temporary arrays, views among them. The memo holds `obj`, so + # its id is not reused while the memo holds its ID: the buffer in + # an ID keeps alive the base of a view, not the view. + entry = self._ids.get(id(obj)) + if entry is None: + entry = self._ids[id(obj)] = (obj, self._array_id(obj)) + self.wrote_numpy = True + return entry[1] if isinstance(obj, numpy.generic): - kind = type(obj) - name = self._scalar_names.get(kind) + dtype = obj.dtype + name = self._scalar_names.get(dtype) 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) + name = self._scalar_names[dtype] = _scalar_name(obj) + self.wrote_numpy = True + return ('scalar', name, _scalar_value(numpy, obj)) return None + def _array_id(self, obj): + """Return the persistent ID of the array `obj`, or raise TypeError.""" + # pylint: disable=protected-access,unidiomatic-typecheck + numpy, ma = self._numpy, self._ma + kind = type(obj) + + if kind is numpy.memmap: + # A memory map is stored as the plain array of its data. + obj, kind = obj.view(numpy.ndarray), numpy.ndarray + + if kind is numpy.ndarray: + if obj.dtype.kind == 'O': + self.wrote_objects = True + shape = tuple(int(size) for size in obj.shape) + return ('objects', shape, list(obj.flat)) + if self._c_order and not obj.flags.c_contiguous: + obj = obj.copy(order='C') + return ('npy',) + _pack_array(numpy, obj) + + if kind is numpy.recarray: + return ('records', obj.view(numpy.ndarray)) + + if ma is not None and kind is ma.MaskedArray: + mask = ma.getmask(obj) + # numpy keeps None until a fill value is set, and its default + # depends on the dtype, so None is stored as None. Checked on + # numpy 1.26 and 2.5. + fill = obj._fill_value + return ( + 'masked', + obj.data, + None if mask is ma.nomask else numpy.asarray(mask), + None if fill is None else numpy.asarray(obj.fill_value), + bool(obj.hardmask), + ) + + name = '%s.%s' % (kind.__module__, kind.__qualname__) + if ma is not None and isinstance(obj, ma.MaskedArray): + raise TypeError( + 'SafeDisk cannot store %s, a subclass of' + ' numpy.ma.MaskedArray; store a plain masked array instead' + % (name,) + ) + raise TypeError( + 'SafeDisk cannot store a subclass of ndarray: %s; store' + ' numpy.asarray() of it instead' % (name,) + ) + class _ArrayUnpickler(SafeUnpickler): """Unpickler that reads back what _ArrayPickler wrote. @@ -252,37 +515,179 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._arrays = {} self._strings = {} + self._built = {} 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)) + kind = pid[0] if type(pid) is tuple and pid else None + if type(kind) is not str: + kind = None + + if kind == 'scalar' and len(pid) == 3: + 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 kind == 'npy' and len(pid) == 3: + 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: - 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] + array = _unpack_array(header, payload) + entry = (header, payload, array) + self._arrays[key] = entry + return entry[2] + + build = _BUILT_KINDS.get(kind) + if build is not None and len(pid) == build[1]: + # The pickle memoizes an ID, so an object named twice is built + # once and stays shared. Holding the ID keeps its id in use. + entry = self._built.get(id(pid)) + if entry is None: + entry = self._built[id(pid)] = (pid, build[0](*pid[1:])) + return entry[1] + raise pickle.UnpicklingError('SafeDisk refuses a persistent ID') +def plain_key(numpy, key): + """Return `key` with numpy scalars in it turned into Python ones. + + numpy booleans, integers, floats, complex numbers, strings and bytes + become the Python value ``item()`` gives, also inside tuples, lists, + dicts, sets and frozensets, so ``int64(1)`` and ``1``, which compare + equal, name the same entry. Other keys are returned unchanged. + + """ + # pylint: disable=unidiomatic-typecheck + kind = type(key) + + if kind in _PLAIN_TYPES: + return key + if isinstance(key, numpy.generic): + return key.item() if key.dtype.kind in 'biufcUS' else key + + if kind is tuple or kind is list: + if all(type(item) in _PLAIN_TYPES for item in key): + return key + items = [plain_key(numpy, item) for item in key] + changed = any(new is not old for new, old in zip(items, key)) + elif kind is dict: + items = [ + (plain_key(numpy, name), plain_key(numpy, value)) + for name, value in key.items() + ] + changed = any( + new[0] is not old[0] or new[1] is not old[1] + for new, old in zip(items, key.items()) + ) + elif kind is set or kind is frozenset: + items = [plain_key(numpy, item) for item in key] + changed = any(new is not old for new, old in zip(items, key)) + else: + return key + + return kind(items) if changed else key + + +_PLAIN_TYPES = frozenset((bytes, str, int, float, bool, type(None))) + + +def _dump(disk, numpy, obj, what, c_order=False): + """Return the pickle of `obj` and whether it holds numpy data. + + Raise TypeError unless the reader of `disk` would accept the pickle. + + """ + # pylint: disable=protected-access + buffer = io.BytesIO() + protocol = max(disk.pickle_protocol, _MIN_PROTOCOL) + pickler = _ArrayPickler(buffer, protocol, numpy, c_order) + + try: + pickler.dump(obj) + except (pickle.PicklingError, AttributeError, RecursionError) as error: + raise TypeError( + 'SafeDisk cannot store %s: %s' % (what, 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, obj, what) + + # An array of objects can hold itself, which pickles but cannot be + # built back, so a pickle with one is read back once before it is kept. + if pickler.wrote_objects: + try: + disk._unpickler(io.BytesIO(data)).load() + except pickle.UnpicklingError as error: + raise TypeError( + 'SafeDisk cannot store %s: %s' % (what, error) + ) from None + + return data, pickler.wrote_numpy + + +def put_key(disk, numpy, key): + """Return the Cache table key fields of `key`, or None to store it as + a plain SafeDisk does, which `key` then passes. + + A key without numpy is pickled as a plain SafeDisk pickles it, so it + keeps the bytes earlier versions wrote. A key holding numpy data is + pickled as a value is, with every array in C order. + + """ + # pylint: disable=protected-access + try: + disk._check(key, 'key') + except TypeError: + pass + else: + return None + + if not _holds_numpy(numpy, key): + try: + data = pickle.dumps(key, protocol=disk.pickle_protocol) + except (pickle.PicklingError, AttributeError, RecursionError): + data = None + if data is not None: + names = _pickle_globals(data) + if names is not None and not any( + _in_modules(module, ('numpy',)) for module, _ in names + ): + disk._check_pickle(data, key, 'key', names) + return sqlite3.Binary(pickletools.optimize(data)), False + + data, _ = _dump(disk, numpy, key, 'key', c_order=True) + return sqlite3.Binary(pickletools.optimize(data)), False + + +def _holds_numpy(numpy, key): + """Return whether `key`, or an item of it if it is a tuple or a list, + is a numpy array or scalar, which spares pickling it twice.""" + # pylint: disable=unidiomatic-typecheck + types = (numpy.ndarray, numpy.generic) + if isinstance(key, types): + return True + return type(key) in (tuple, list) and any( + isinstance(item, types) for item in key + ) + + def unpickler(disk, file): """Return the unpickler of `disk`, a SafeDisk with ``safe_numpy`` set.""" return _ArrayUnpickler( @@ -308,22 +713,5 @@ def store(disk, numpy, value, key): 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') - + data, _ = _dump(disk, numpy, value, 'value') return disk._store_pickle(data, value, key) diff --git a/diskcache_community/arrays.pyi b/diskcache_community/arrays.pyi index 84a46f0..b7ce043 100644 --- a/diskcache_community/arrays.pyi +++ b/diskcache_community/arrays.pyi @@ -3,6 +3,10 @@ from typing import IO, Any from diskcache_community.core import SafeDisk, SafeUnpickler def unpickler(disk: SafeDisk, file: IO[bytes]) -> SafeUnpickler: ... +def plain_key(numpy: Any, key: Any) -> Any: ... +def put_key( + disk: SafeDisk, numpy: Any, key: Any +) -> tuple[Any, bool] | None: ... 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 b2dcb86..98d2cac 100644 --- a/diskcache_community/core.py +++ b/diskcache_community/core.py @@ -955,13 +955,24 @@ def with_numpy(cls): The returned class writes arrays through pickle's persistent IDs in numpy's ``.npy`` format with ``allow_pickle=False``, anywhere in a - value except in keys, and checks each header before it reads an - array back. (Plain :class:`SafeDisk` refuses numpy pickles, because - unpickling an array is not memory-safe.) numpy scalars keep their - type, and arrays of ``datetime64`` and ``timedelta64`` are stored too. - 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. + value, and checks each header before it reads an array back. (Plain + :class:`SafeDisk` refuses numpy pickles, because unpickling an array + is not memory-safe.) numpy scalars in a value keep their type, + ``datetime64`` and ``timedelta64`` scalars their unit. Arrays of + objects are stored element by element, each element checked like any + value. Record arrays and masked arrays keep their type, and a memory + map is stored as the plain array of its data. Other ``ndarray`` + subclasses, such as ``numpy.matrix``, arrays with object fields and + other numpy scalars raise :exc:`TypeError` when written. It combines + with :meth:`with_types` and :meth:`with_options` in any order. + + In a key, numpy booleans, integers, floats, complex numbers, strings + and bytes, also inside tuples, lists, dicts, sets and frozensets, + become the Python values ``item()`` gives, so ``int64(1)`` and ``1`` + name the same entry. Arrays and the other + numpy objects above are stored in keys as in values, with every array + in C order, so equal arrays of the same dtype and shape name the same + entry. A key without numpy is stored as a plain SafeDisk stores it. 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, @@ -1015,9 +1026,14 @@ def _check_object(self, obj, what): self._check_pickle(data, obj, what) - def _check_pickle(self, data, obj, what): - """Raise TypeError unless the reader would accept the pickle `data`.""" - names = _pickle_globals(data) + def _check_pickle(self, data, obj, what, names=UNKNOWN): + """Raise TypeError unless the reader would accept the pickle `data`. + + `names` are the names `data` looks up, if the caller already has them. + + """ + if names is UNKNOWN: + names = _pickle_globals(data) if names is None: try: @@ -1041,15 +1057,20 @@ def _check_pickle(self, data, obj, what): type(obj).__qualname__, module, name, - self._refusal_hint(module, name), + self._refusal_hint(module, name, what), ) ) from None - def _refusal_hint(self, module, name): + def _refusal_hint(self, module, name, what): """Return how to store what pickles the refused `module`.`name`.""" if _in_modules(module, ('numpy',)): if self.safe_numpy: - return '; with_numpy() stores only arrays and numpy scalars' + hint = '; with_numpy() stores arrays and numpy scalars, not other numpy objects' + # pylint: disable-next=protected-access + resolved = SafeUnpickler._resolve(module, name) + if what == 'key' and isinstance(resolved, type): + hint += ', such as the numpy classes that memoize(typed=True) puts in keys' + return hint return '; store arrays with SafeDisk.with_numpy()' for unsafe in _UNSAFE_MODULES: @@ -1143,15 +1164,20 @@ def _confine(self, filename): def put(self, key): # pylint: disable=unidiomatic-typecheck - if type(key) not in (bytes, str, int, float): + if type(key) in (bytes, str, int, float): + return super().put(key) + numpy = sys.modules.get('numpy') if self.safe_numpy else None + if numpy is None: self._check_object(key, 'key') - return super().put(key) + return super().put(key) + key = arrays.plain_key(numpy, key) + fields = arrays.put_key(self, numpy, key) + return super().put(key) if fields is None else fields def get(self, key, raw): if raw: return super().get(key, raw) - # Keys exclude numpy, so a key row gets no persistent IDs. - return self._plain_unpickler(io.BytesIO(key)).load() + return self._unpickler(io.BytesIO(key)).load() def store(self, value, read, key=UNKNOWN): # pylint: disable=unidiomatic-typecheck diff --git a/docs/conf.py b/docs/conf.py index b0901be..02262ca 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -20,8 +20,8 @@ # -- Project information ----------------------------------------------------- project = 'DiskCache' -copyright = '2016-2023 Grant Jenks, 2026 diskcache-community contributors' -author = 'Grant Jenks and diskcache-community contributors' +copyright = '2026 August V.' +author = 'August V.' # The full version, including alpha/beta/rc tags release = diskcache.__version__ @@ -42,7 +42,7 @@ ogp_site_url = 'https://diskcache-community.readthedocs.io/en/latest/' ogp_site_name = 'diskcache-community' ogp_image = 'https://raw.githubusercontent.com/secondwindoss/diskcache-community/master/.github/social-preview.png' -ogp_image_alt = 'diskcache-community: the maintained fork of python-diskcache' +ogp_image_alt = 'diskcache-community' ogp_description_length = 200 # Add any paths that contain templates here, relative to this directory. diff --git a/docs/index.rst b/docs/index.rst index d0510b9..25aa48b 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -8,16 +8,14 @@ processes, evicts by LRU, LFU and other policies, and ships ``Deque`` and ``Index`` containers plus recipes for memoization, locks and throttling. This documentation describes `diskcache-community`_, a community fork of -`python-diskcache`_. The fork is published on PyPI as ``diskcache-community`` -and imported as ``diskcache_community``, so it installs next to the upstream -``diskcache`` package. It keeps the public API, the ``cache.db`` schema and +DiskCache, published on PyPI as ``diskcache-community`` and imported as +``diskcache_community``. It keeps the public API, the ``cache.db`` schema and the file layout of the last upstream release, 5.6.3, so existing cache directories keep working, and it carries the fix path for CVE-2025-69872 through the opt-in :class:`SafeDisk `. The README in the repository explains why the fork exists and how to migrate. -Install it with pip and import ``diskcache_community`` where you imported -``diskcache``:: +Install it with pip and import ``diskcache_community``: $ pip install diskcache-community @@ -45,4 +43,3 @@ stores, and the fork with diskcache 5.6.3. development .. _`diskcache-community`: https://github.com/secondwindoss/diskcache-community -.. _`python-diskcache`: https://github.com/grantjenks/python-diskcache diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 0d3516c..2969127 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -1095,16 +1095,42 @@ in numpy's ``.npy`` format: * 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 `. +* numpy booleans, integers, floats, complex numbers, strings and bytes read + back as the same numpy type with the same value, including NaN. + ``datetime64`` and ``timedelta64`` scalars read back with the same unit, + including ``NaT``; a unit with a multiple of more than nine digits, such as + ``1000000000s``, raises :exc:`TypeError`. +* An array of objects (``dtype=object``), such as a pandas text column from + ``to_numpy()``, is stored element by element and reads back in C order. + Each element is checked like any value, so an element of a class outside + the allowlist raises :exc:`TypeError`. A value holding an array of objects + is read back once before it is stored, which makes writing it slower. +* Record arrays (``numpy.recarray``) and masked arrays + (``numpy.ma.MaskedArray``, with their mask, fill value and hard mask, but + not other state numpy keeps on them) read back with their type. + ``DataFrame.to_records()`` gives a record array, but one with object fields + for a frame with text columns, which is refused: store the columns instead. + A memory map (``numpy.memmap``) is stored as the plain array of its data and + reads back as a plain array. +* Other subclasses of ``ndarray``, such as ``numpy.matrix``, arrays with + object fields, other numpy scalars and arrays that the reader would refuse, + such as a dtype with no width, raise :exc:`TypeError` when written. + +Keys can hold numpy data too. numpy booleans, integers, floats, complex +numbers, strings and bytes, also inside tuples, lists, dicts, sets and +frozensets, become the Python values ``item()`` gives, so ``int64(1)`` and +``1``, which compare equal, name the same entry, and a function decorated with +:meth:`memoize ` hits for either argument. +Arrays and the other numpy objects above are stored in a key as in a value, +with every array in C order: an equal array of the same dtype and shape names +the same entry, in either memory order, while the same numbers in another +dtype or shape name another one, and so do equal times in another unit. As +with any pickled key, a key that holds one object twice, such as ``(a, a)``, +names another entry than one holding two equal objects, ``(a, a.copy())``. A +key without numpy is stored exactly as a plain :class:`SafeDisk +` stores it. numpy arguments with +``memoize(typed=True)`` raise :exc:`TypeError`, since the key then holds +numpy's classes. Everything else behaves as in :class:`SafeDisk `, and ``with_options(on_error='miss')`` works @@ -1305,7 +1331,6 @@ files. It does not account the size of directories themselves or other filesystem metadata. If directory count or size is a concern then consider implementing an alternative :class:`Disk `. -.. _`#54`: https://github.com/grantjenks/python-diskcache/issues/54 .. _`hash protocol`: https://docs.python.org/library/functions.html#hash .. _`not recommended`: https://www.sqlite.org/faq.html#q5 .. _`performs poorly`: https://www.pythonanywhere.com/forums/topic/1847/ diff --git a/pyproject.toml b/pyproject.toml index a669dea..b395584 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,13 +5,13 @@ build-backend = "setuptools.build_meta" [project] name = "diskcache-community" dynamic = ["version"] -description = "Disk Cache -- Disk and file backed persistent cache." +description = "Disk Cache - Disk and file backed persistent cache." readme = "README.md" requires-python = ">=3.8" license = "Apache-2.0" -license-files = ["LICENSE", "NOTICE"] -authors = [{ name = "Grant Jenks" }] -maintainers = [{ name = "diskcache-community contributors" }] +license-files = ["LICENSE"] +authors = [{ name = "August V." }] +maintainers = [{ name = "August V." }] classifiers = [ "Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", @@ -33,7 +33,6 @@ Homepage = "https://github.com/secondwindoss/diskcache-community" Documentation = "https://diskcache-community.readthedocs.io/" Source = "https://github.com/secondwindoss/diskcache-community" Tracker = "https://github.com/secondwindoss/diskcache-community/issues" -Upstream = "https://github.com/grantjenks/python-diskcache" [tool.setuptools] packages = ["diskcache_community"] diff --git a/tests/benchmark_core.py b/tests/benchmark_core.py index 2e9f1cc..2c5feeb 100644 --- a/tests/benchmark_core.py +++ b/tests/benchmark_core.py @@ -1,10 +1,3 @@ -"""Benchmark diskcache_community.Cache - -$ export PYTHONPATH=/Users/grantj/repos/python-diskcache -$ python tests/benchmark_core.py -p 1 > tests/timings_core_p1.txt -$ python tests/benchmark_core.py -p 8 > tests/timings_core_p8.txt -""" - import collections as co import multiprocessing as mp import os diff --git a/tests/benchmark_djangocache.py b/tests/benchmark_djangocache.py index f48cbfd..c00e0e7 100644 --- a/tests/benchmark_djangocache.py +++ b/tests/benchmark_djangocache.py @@ -1,9 +1,3 @@ -"""Benchmark diskcache_community.DjangoCache - -$ export PYTHONPATH=/Users/grantj/repos/python-diskcache -$ python tests/benchmark_djangocache.py > tests/timings_djangocache.txt -""" - import collections as co import multiprocessing as mp import os diff --git a/tests/plot.py b/tests/plot.py index fcac0bc..ec0a1e4 100644 --- a/tests/plot.py +++ b/tests/plot.py @@ -1,6 +1,5 @@ """Plot Benchmarks for docs -$ export PYTHONPATH=/Users/grantj/repos/python-diskcache $ python tests/plot.py --show tests/timings_core_p1.txt """ diff --git a/tests/test_safedisk_numpy.py b/tests/test_safedisk_numpy.py index e4ba060..99b83ff 100644 --- a/tests/test_safedisk_numpy.py +++ b/tests/test_safedisk_numpy.py @@ -115,6 +115,16 @@ def persistent_id(self, obj): return obj.pid if isinstance(obj, Pid) else None +class Call: + """Pickles as a call of `func` with `args`, as a crafted row can.""" + + def __init__(self, func, *args): + self.func, self.args = func, args + + def __reduce__(self): + return self.func, self.args + + def forge(value): buffer = io.BytesIO() Forger(buffer, protocol=5).dump(value) @@ -509,20 +519,50 @@ def test_scalars_in_the_other_safe_containers(directory, make): assert type(item) is np.int64 +TIME_SCALARS = [ + np.datetime64('2026-01-01T00:00:00.000000000'), + np.datetime64('2026-11-05'), + np.datetime64('2026-11', 'M'), + np.datetime64(7, '25s'), + np.datetime64('NaT', 'ns'), + np.timedelta64(5, 'ns'), + np.timedelta64(-3, 'W'), + np.timedelta64('NaT', 'D'), + np.timedelta64(4, 'h'), +] + + +@pytest.mark.parametrize('scalar', TIME_SCALARS, ids=repr) +def test_time_scalars_keep_their_unit(directory, scalar): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = [scalar] + (result,) = cache['key'] + cache[scalar] = 'by time' + assert cache[scalar] == 'by time' + assert type(result) is type(scalar) + assert result.dtype == scalar.dtype + assert result.view(np.int64) == scalar.view(np.int64) + + @pytest.mark.parametrize( - 'scalar', + 'name, value', [ - np.datetime64('2026-01-01T00:00:00.000000000'), - np.timedelta64(5, 'ns'), + ('datetime64[D]', 'text'), + ('datetime64[D]', 2**63), + ('datetime64[D]', 1.5), + ('datetime64[parsec]', 1), + ('datetime64[0s]', 1), + ('datetime64[D] ', 1), + ('timedelta64[1234567890s]', 1), + ('datetime64,f8', 1), ], - ids=lambda value: type(value).__name__, ) -def test_time_scalars_are_refused(directory, scalar): +def test_forged_time_scalar_is_refused(directory, name, value): + write_row(directory, persid(('scalar', name, value))) + with dc.Cache(directory, disk=ArrayDisk) as cache: - with pytest.raises(TypeError): - cache['key'] = scalar - with pytest.raises(TypeError): - cache['key'] = [scalar] + with pytest.raises(pickle.UnpicklingError): + cache['key'] def test_long_double_reads_back_or_is_refused(directory): @@ -539,11 +579,11 @@ def test_long_double_reads_back_or_is_refused(directory): @pytest.mark.parametrize( 'array', [ - np.array([object()], dtype=object), + np.array([Outside()], dtype=object), np.array([(1, None)], dtype=[('a', 'i4'), ('b', 'O')]), np.zeros(1, dtype=[('a', 'O', (2,))]), ], - ids=['object', 'field', 'subarray'], + ids=['outside-element', 'field', 'subarray'], ) def test_object_arrays_are_refused(directory, array): with dc.Cache(directory, disk=ArrayDisk) as cache: @@ -576,30 +616,368 @@ def test_arrays_the_reader_would_refuse_fail_at_write(directory, array): assert len(cache) == 0 +class Matrixish(np.ndarray): + """A subclass of ndarray from outside numpy.""" + + @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'], + [lambda: np.matrix([1, 2]), lambda: np.zeros(2).view(Matrixish)], + ids=['matrix', 'outside'], ) -def test_subclasses_are_refused(directory, make): +def test_other_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'): + with pytest.raises(TypeError, match='subclass.*numpy.asarray'): cache['key'] = make() with pytest.raises(TypeError, match='subclass'): cache['key'] = [make()] + with pytest.raises(TypeError, match='subclass'): + cache[make()] = 'value' -def test_numpy_in_a_key_is_refused(directory): +# Object arrays + + +@pytest.mark.parametrize( + 'array', + [ + np.array(['a', 1, None, 2.5, b'x', (1, 2)], dtype=object), + np.array([[1, 'a'], [None, 2]], dtype=object), + np.array(7, dtype=object), + np.empty((0, 3), dtype=object), + np.asfortranarray(np.array([[1, 2], [3, 4]], dtype=object)), + ], + ids=['mixed', '2-d', '0-d', 'empty', 'fortran'], +) +def test_object_array_round_trip(directory, array): with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = array + result = cache['key'] + assert type(result) is np.ndarray + assert result.dtype == object + assert result.shape == array.shape + assert result.tolist() == array.tolist() + + +def test_object_array_holds_lists_arrays_and_shared_items(directory): + inner = np.arange(3.0) + array = np.empty(4, dtype=object) + array[0] = [1, 2] + array[1] = inner + array[2] = inner + array[3] = {'nested': np.array(['x', None], dtype=object)} + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = array + result = cache['key'] + + assert result[0] == [1, 2] + same(result[1], inner) + assert result[1] is result[2] + assert result[3]['nested'].tolist() == ['x', None] + + +@pytest.mark.parametrize( + 'args', + [ + ((3,), [1, 2]), + ([2], [1, 2]), + ((2, -1), [1, 2]), + ((True,), [1]), + ((1,) * 65, [1]), + ((2,), (1, 2)), + ((10**20, 0), []), + ], + ids=['count', 'list-shape', 'negative', 'bool', 'dims', 'tuple', 'huge'], +) +def test_forged_object_array_is_refused(directory, args): + write_row(directory, persid(('objects',) + args)) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(pickle.UnpicklingError): + cache['key'] + + with dc.Cache(directory, disk=MissDisk) as cache: + assert cache.get('key', 'default') == 'default' + + +@pytest.mark.parametrize( + 'pid', + [ + ('objects', (1,), [1]), + ('records', Pid(array_pid(npy_parts(np.zeros(1))))), + ('masked', Pid(array_pid(npy_parts(np.zeros(1)))), None, None, False), + ], + ids=['objects', 'records', 'masked'], +) +def test_built_kinds_need_with_numpy(directory, pid): + write_row(directory, persid(pid)) + + with dc.Cache(directory, disk=dc.SafeDisk) as cache: + with pytest.raises(pickle.UnpicklingError): + cache['key'] + + with dc.Cache(directory, disk=ArrayDisk) as cache: + assert cache['key'].shape == (1,) + + +@pytest.mark.parametrize( + 'value', + [ + Call(arrays._object_array, (1,), [1]), + Call(arrays._import_numpy), + Pid((Pid(array_pid(npy_parts(np.arange(2.0)))), 1, 2)), + Pid(('objects', (1,))), + Pid(('records', Pid(array_pid(npy_parts(np.zeros(1)))), 1)), + ], + ids=['function', 'module-name', 'array-kind', 'short', 'long'], +) +def test_other_rows_naming_arrays_are_refused(directory, value): + write_row(directory, forge(value)) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(pickle.UnpicklingError): + cache['key'] + + +# Record arrays, masked arrays and memory maps + + +def test_record_array_round_trip(directory): + records = np.rec.array( + [(1, 2.5, 'a'), (2, 3.5, 'b')], + dtype=[('x', 'i4'), ('y', 'f8'), ('name', 'U3')], + ) + plain = np.array([1.0, 2.0]).view(np.recarray) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['records'] = records + cache['plain'] = [plain] + result = cache['records'] + (plain_result,) = cache['plain'] + + assert type(result) is np.recarray + assert result.dtype == records.dtype + assert result.x.tolist() == [1, 2] + assert result.name.tolist() == ['a', 'b'] + assert type(plain_result) is np.recarray + assert plain_result.tolist() == [1.0, 2.0] + + +def test_forged_record_array_is_refused(directory): + write_row(directory, persid(('records', [1, 2]))) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(pickle.UnpicklingError): + cache['key'] + + +@pytest.mark.parametrize( + 'array', + [ + np.ma.masked_array([1.0, 2.0, 3.0], mask=[0, 1, 0], fill_value=-1), + np.ma.masked_array([1, 2]), + np.ma.masked_array([[1, 2], [3, 4]], mask=[[1, 0], [0, 1]]), + np.ma.masked_array(['a', 'b'], mask=[0, 1]), + np.ma.masked_array( + [(1, 2.0), (3, 4.0)], + mask=[(0, 1), (1, 0)], + dtype=[('a', 'i4'), ('b', 'f8')], + ), + np.ma.masked_array([1, 2], mask=[1, 0], hard_mask=True), + ], + ids=['float', 'no-mask', '2-d', 'str', 'structured', 'hard'], +) +def test_masked_array_round_trip(directory, array): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = array + result = cache['key'] + + assert type(result) is np.ma.MaskedArray + assert result.dtype == array.dtype + assert result.data.tobytes() == array.data.tobytes() + assert np.ma.getmaskarray(result).tolist() == ( + np.ma.getmaskarray(array).tolist() + ) + assert result.fill_value == array.fill_value + assert result.hardmask == array.hardmask + if np.ma.getmask(array) is np.ma.nomask: + assert np.ma.getmask(result) is np.ma.nomask + + +@pytest.mark.parametrize( + 'args', + [ + ([1.0], None, np.array(0.0), False), + (np.zeros(2), np.zeros(3, dtype=bool), np.array(0.0), False), + (np.zeros(2), np.zeros(2, dtype='i1'), np.array(0.0), False), + (np.zeros(2), None, np.array('text'), False), + (np.zeros(2), None, np.zeros(1), False), + (np.zeros(2), None, np.array(0.0), 1), + ], + ids=[ + 'data', + 'mask-shape', + 'mask-dtype', + 'fill-text', + 'fill-shape', + 'hard', + ], +) +def test_forged_masked_array_is_refused(directory, args): + data, mask, fill, hard = ( + Pid(array_pid(npy_parts(arg))) if isinstance(arg, np.ndarray) else arg + for arg in args + ) + write_row(directory, persid(('masked', data, mask, fill, hard))) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(pickle.UnpicklingError): + cache['key'] + + +def test_memory_map_is_stored_as_its_data(directory): + path = op.join(directory, 'map.bin') + mapped = np.memmap(path, dtype='f4', mode='w+', shape=(2, 3)) + mapped[:] = np.arange(6).reshape(2, 3) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = {'map': mapped, 'row': mapped[1]} + result = cache['key'] + cache[mapped] = 'by map' + assert cache[np.arange(6, dtype='f4').reshape(2, 3)] == 'by map' + + same(result['map'], np.arange(6, dtype='f4').reshape(2, 3)) + same(result['row'], np.arange(3, 6, dtype='f4')) + del mapped + + +# Arrays in keys + + +def test_array_key_names_the_entry_of_an_equal_array(directory): + key = np.arange(6.0).reshape(2, 3) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache[key] = 'value' + assert cache[key.copy()] == 'value' + assert cache[np.asfortranarray(key)] == 'value' + assert cache[key[::-1][::-1]] == 'value' + assert cache.get(key.astype(np.float32)) is None + assert cache.get(key.reshape(3, 2)) is None + assert cache.get(key + 1) is None + assert len(cache) == 1 + (stored,) = list(cache) + + same(stored, key) + + +def test_array_key_bytes_do_not_depend_on_order_or_flags(directory): + disk = ArrayDisk(directory) + key = np.arange(6.0).reshape(2, 3) + readonly = key.copy() + readonly.flags.writeable = False + expected = disk.put(key) + + assert disk.put(np.asfortranarray(key)) == expected + assert disk.put(readonly) == expected + assert disk.hash(np.asfortranarray(key)) == disk.hash(key) + + +@pytest.mark.parametrize( + 'key', + [ + ('image', np.arange(4, dtype=np.uint8)), + [np.zeros(2), {'x': np.ones(1)}.get('x')], + (np.datetime64('2026-11-05'), 'day'), + np.array(['a', None, 3], dtype=object), + np.rec.array([(1, 2.0)], dtype=[('a', 'i4'), ('b', 'f8')]), + np.ma.masked_array([1, 2], mask=[0, 1]), + ], + ids=['tuple', 'list', 'time', 'objects', 'records', 'masked'], +) +def test_numpy_keys_round_trip(directory, key): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache[key] = 'value' + assert cache[key] == 'value' + (stored,) = list(cache) + assert repr(stored) == repr(key) + + +def test_keys_without_numpy_keep_their_bytes(directory): + for key in [('a', 1), Pair(1, 2), frozenset({1})]: + assert ArrayDisk.with_types(Pair)(directory).put(key) == ( + dc.SafeDisk.with_types(Pair)(directory).put(key) + ) + + +def test_memoize_with_array_arguments(directory): + calls = [] + + with dc.Cache(directory, disk=ArrayDisk) as cache: + + @cache.memoize() + def total(values): + calls.append(1) + return float(values.sum()) + + assert total(np.arange(4.0)) == 6.0 + assert total(np.arange(4.0)) == 6.0 + assert total(np.arange(5.0)) == 10.0 + + assert len(calls) == 2 + + +def test_fanout_cache_with_array_keys(directory): + with dc.FanoutCache(directory, shards=4, disk=ArrayDisk) as cache: + for size in range(8): + cache[np.arange(float(size))] = size + assert [cache[np.arange(float(size))] for size in range(8)] == list( + range(8) + ) + + +@pytest.mark.parametrize( + 'numpy_key, plain', + [ + (np.int64(1), 1), + (np.uint8(200), 200), + (np.float64(2.5), 2.5), + (np.float32(0.5), 0.5), + (np.bool_(True), True), + (np.str_('text'), 'text'), + (np.bytes_(b'data'), b'data'), + ((np.int64(8), np.int64(3)), (8, 3)), + (('name', [np.int32(1), (np.str_('a'),)]), ('name', [1, ('a',)])), + ], +) +def test_numpy_scalar_in_a_key_names_the_plain_entry( + directory, numpy_key, plain +): + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache[numpy_key] = 'value' + assert cache[plain] == 'value' + assert numpy_key in cache + assert len(cache) == 1 + (stored,) = list(cache) + assert stored == plain + assert type(stored) is type(plain) + cache[plain] = 'other' + assert cache[numpy_key] == 'other' + assert len(cache) == 1 + + +def test_numpy_scalar_key_is_refused_without_numpy_support(directory): + with dc.Cache(directory, disk=dc.SafeDisk) 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_plain_key_leaves_other_keys_alone(): + key = ('a', 1, [2.5, b'x'], None) + assert arrays.plain_key(np, key) is key + assert arrays.plain_key(np, (np.int8(1),)) == (1,) def test_memoize(directory): @@ -631,10 +1009,27 @@ def compute(): assert np.isnan(first) and np.isnan(second) -def test_memoize_scalar_argument_is_refused(directory): +def test_memoize_scalar_argument_shares_the_plain_entry(directory): + calls = [] + with dc.Cache(directory, disk=ArrayDisk) as cache: @cache.memoize() + def compute(value): + calls.append(value) + return value * 2 + + assert compute(np.int64(3)) == 6 + assert compute(3) == 6 + assert compute(np.int32(3)) == 6 + + assert calls == [3] + + +def test_memoize_typed_scalar_argument_is_refused(directory): + with dc.Cache(directory, disk=ArrayDisk) as cache: + + @cache.memoize(typed=True) def compute(value): return value @@ -935,14 +1330,19 @@ def test_with_numpy_refusal_of_other_numpy_objects(directory): with pytest.raises(TypeError) as info: cache['key'] = {'dtype': np.dtype('f8')} assert str(info.value).endswith( - 'with_numpy() stores only arrays and numpy scalars' + 'with_numpy() stores arrays and numpy scalars, not other numpy' + ' objects' ) -def test_object_array_refusal_mentions_dataframes(directory): +def test_object_field_refusal_mentions_dataframes(directory): with dc.Cache(directory, disk=ArrayDisk) as cache: - with pytest.raises(TypeError, match='DataFrame .* records'): - cache['key'] = np.array(['a', None], dtype=object) + with pytest.raises( + TypeError, match='object fields.*to_records.*columns' + ): + cache['key'] = np.array( + [(1, None)], dtype=[('a', 'i4'), ('b', 'O')] + ) def test_safejsondisk_refuses_with_numpy(): @@ -1210,22 +1610,39 @@ def test_forged_persistent_id_is_refused(directory, pid): assert len(cache) == 0 -@pytest.mark.parametrize( - 'pid', - [array_pid(npy_parts(np.arange(2.0))), ('scalar', 'int64', 1)], - ids=['array', 'scalar'], -) -def test_persistent_id_in_a_key_row_is_refused(directory, pid): +def write_key_row(directory, data): + """Store the pickle `data` in the key 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 key = ? WHERE raw = 0', - (sqlite3.Binary(persid(pid)),), + 'UPDATE Cache SET key = ? WHERE raw = 0', (sqlite3.Binary(data),) ) con.commit() con.close() + +@pytest.mark.parametrize( + 'pid', + [array_pid(npy_parts(np.arange(2.0))), ('scalar', 'int64', 1)], + ids=['array', 'scalar'], +) +def test_persistent_id_in_a_key_row(directory, pid): + write_key_row(directory, persid(pid)) + + with dc.Cache(directory, disk=dc.SafeDisk) as cache: + with pytest.raises(pickle.UnpicklingError): + list(cache) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + (key,) = list(cache) + assert repr(key) in (repr(np.arange(2.0)), repr(np.int64(1))) + + +@pytest.mark.parametrize('pid', FORGED_IDS.values(), ids=FORGED_IDS.keys()) +def test_forged_persistent_id_in_a_key_row_is_refused(directory, pid): + write_key_row(directory, persid(pid)) + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(pickle.UnpicklingError): list(cache) @@ -1436,3 +1853,304 @@ def test_deep_nesting_reads_without_recursing(directory): result = cache.get('key', 'default') assert type(result) is list + + +def test_header_is_parsed_once_and_still_checked_against_the_data( + directory, monkeypatch +): + arrays._parsed_headers.clear() + arrays._written_headers.clear() + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['a'] = np.arange(4, dtype=np.float32) + parsed = [] + original = arrays._parse_header + monkeypatch.setattr( + arrays, + '_parse_header', + lambda *args: parsed.append(1) or original(*args), + ) + cache['b'] = np.arange(4, 8, dtype=np.float32) + same(cache['a'], np.arange(4, dtype=np.float32)) + same(cache['b'], np.arange(4, 8, dtype=np.float32)) + assert parsed == [] + + header, _ = arrays._pack_array(np, np.arange(4, dtype=np.float32)) + with pytest.raises(pickle.UnpicklingError): + arrays._unpack_array(header, bytearray(12)) + with pytest.raises(pickle.UnpicklingError): + arrays._unpack_array(header, bytearray(20)) + + +def test_written_header_matches_numpy(directory): + arrays._written_headers.clear() + for array in ( + np.zeros((3, 2), dtype='>i2'), + np.asfortranarray(np.zeros((3, 2))), + np.zeros(0, dtype='M8[ns]'), + np.zeros((), dtype='U3'), + ): + for _ in range(2): + header, _ = arrays._pack_array(np, array) + buffer = io.BytesIO() + fields = np.lib.format.header_data_from_array_1_0(array) + np.lib.format.write_array_header_1_0(buffer, fields) + assert header == buffer.getvalue() + + +def test_structured_dtype_headers_are_not_kept(directory): + arrays._parsed_headers.clear() + arrays._written_headers.clear() + dtype = np.dtype([('x', ' Date: Thu, 5 Nov 2026 22:31:37 -0800 Subject: [PATCH 2/2] Name upstream issue 54 in plain text in the tutorial, whose link target was removed. --- docs/tutorial.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 2969127..9742a84 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -1294,9 +1294,9 @@ both keys and values. Unfortunately, pickling produces inconsistencies sometimes when applied to container data types like tuples. Two equal tuples may serialize to different bytes objects using pickle. The likelihood of differences is reduced by using `pickletools.optimize` but still -inconsistencies occur (`#54`_). The inconsistent serialized pickle values is -particularly problematic when applied to the key in the cache. Consider using -an alternative Disk type, like :class:`JSONDisk +inconsistencies occur (upstream issue 54). The inconsistent serialized pickle +values is particularly problematic when applied to the key in the cache. +Consider using an alternative Disk type, like :class:`JSONDisk `, for consistent serialization of keys. SQLite is used to synchronize database access between threads and processes and