From 0d8d1efb49c10bb47ac6aa8a051ca9fa41b0db63 Mon Sep 17 00:00:00 2001 From: August Vi Date: Sun, 1 Nov 2026 21:16:28 -0800 Subject: [PATCH 01/10] Add SafeArrayDisk, a SafeDisk that stores numpy arrays as .npy bytes. --- CHANGELOG.md | 2 + README.md | 4 +- diskcache_community/__init__.py | 2 + diskcache_community/__init__.pyi | 2 + diskcache_community/arrays.py | 183 +++++++++++ diskcache_community/arrays.pyi | 3 + diskcache_community/core.py | 5 +- docs/api.rst | 10 + docs/fix-cve-2025-69872.rst | 16 +- docs/tutorial.rst | 46 ++- tests/test_safearraydisk.py | 543 +++++++++++++++++++++++++++++++ 11 files changed, 800 insertions(+), 16 deletions(-) create mode 100644 diskcache_community/arrays.py create mode 100644 diskcache_community/arrays.pyi create mode 100644 tests/test_safearraydisk.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9d2a2f8..728e9c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ 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). +- `SafeArrayDisk` is a `SafeDisk` that stores numpy arrays in the `.npy` format instead of pickling them, on their own or inside a `dict`, `list` or `tuple`, and stores numpy scalars as Python values. Array headers are checked against the data before reading, so a forged shape is refused rather than raising `MemoryError`. Object arrays and `ndarray` subclasses raise `TypeError` at write time. numpy stays optional and refused by `SafeDisk` (PR #34). + ### Fixes - `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). diff --git a/README.md b/README.md index 8be58df..a20f507 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`, `SafeJSONDisk` and `SafeArrayDisk` for caches whose + directory other users can write to - 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/diskcache_community/__init__.py b/diskcache_community/__init__.py index d97178a..4c7f1b1 100644 --- a/diskcache_community/__init__.py +++ b/diskcache_community/__init__.py @@ -5,6 +5,7 @@ The :doc:`tutorial` provides a helpful walkthrough of most methods. """ +from .arrays import SafeArrayDisk from .core import ( DEFAULT_SETTINGS, ENOVAL, @@ -50,6 +51,7 @@ 'Lock', 'RLock', 'ReadOnlyError', + 'SafeArrayDisk', 'SafeDisk', 'SafeJSONDisk', 'Timeout', diff --git a/diskcache_community/__init__.pyi b/diskcache_community/__init__.pyi index b3f34e5..977da3f 100644 --- a/diskcache_community/__init__.pyi +++ b/diskcache_community/__init__.pyi @@ -1,3 +1,4 @@ +from diskcache_community.arrays import SafeArrayDisk as SafeArrayDisk from diskcache_community.core import DEFAULT_SETTINGS as DEFAULT_SETTINGS from diskcache_community.core import ENOVAL as ENOVAL from diskcache_community.core import EVICTION_POLICY as EVICTION_POLICY @@ -40,6 +41,7 @@ __all__ = [ 'Lock', 'RLock', 'ReadOnlyError', + 'SafeArrayDisk', 'SafeDisk', 'SafeJSONDisk', 'Timeout', diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py new file mode 100644 index 0000000..3ff1843 --- /dev/null +++ b/diskcache_community/arrays.py @@ -0,0 +1,183 @@ +"""Disk that stores numpy arrays without pickling them.""" + +import io +import math +import pickle +import sys + +from .core import UNKNOWN, SafeDisk + +_NPY_KEY = '__diskcache_community_npy__' + + +def _contains(value, match): + """Return True if `match` holds for `value` or anything nested in it.""" + seen = set() + stack = [value] + + while stack: + item = stack.pop() + if match(item): + return True + type_item = type(item) + if type_item in (dict, list, tuple): + if id(item) in seen: + continue + seen.add(id(item)) + if type_item is dict: + stack.extend(item.keys()) + stack.extend(item.values()) + else: + stack.extend(item) + + return False + + +def _rebuild(value, convert): + """Return `value` with `convert` applied through dicts, lists and tuples. + + `convert` returns its argument unchanged for anything it leaves alone. + + """ + converted = convert(value) + if converted is not value: + return converted + type_value = type(value) + if type_value is dict: + return { + _rebuild(key, convert): _rebuild(item, convert) + for key, item in value.items() + } + if type_value is list: + return [_rebuild(item, convert) for item in value] + if type_value is tuple: + return tuple(_rebuild(item, convert) for item in value) + return value + + +def _encoded_array(item): + # pylint: disable=unidiomatic-typecheck + return type(item) is dict and len(item) == 1 and _NPY_KEY in item + + +def _pack_array(numpy, array): + """Return `array` as a dict that holds its ``.npy`` bytes.""" + if array.dtype.hasobject: + raise TypeError( + 'SafeArrayDisk cannot store an array of dtype %s' % (array.dtype,) + ) + buffer = io.BytesIO() + numpy.save(buffer, array, allow_pickle=False) + data = buffer.getvalue() + if data[6:8] not in (b'\x01\x00', b'\x02\x00'): + raise TypeError( + 'SafeArrayDisk cannot store an array of dtype %s' % (array.dtype,) + ) + return {_NPY_KEY: data} + + +def _unpack_array(data): + """Return the array in ``.npy`` bytes that may have been tampered with. + + The sizes in the header are checked against the data first, because a + forged shape would otherwise raise :exc:`MemoryError`, which + ``on_error='miss'`` does not treat as an unreadable entry. + + """ + # pylint: disable=import-error,import-outside-toplevel,no-member + import numpy + + reader = io.BytesIO(data) + version = numpy.lib.format.read_magic(reader) + if version == (1, 0): + header = numpy.lib.format.read_array_header_1_0(reader) + elif version == (2, 0): + header = numpy.lib.format.read_array_header_2_0(reader) + else: + raise pickle.UnpicklingError('SafeArrayDisk found no array data') + + shape, _, dtype = header + size = dtype.itemsize * math.prod(shape) + + if size != len(data) - reader.tell(): + raise pickle.UnpicklingError('SafeArrayDisk refuses the array header') + + reader.seek(0) + return numpy.lib.format.read_array(reader, allow_pickle=False) + + +class SafeArrayDisk(SafeDisk): + """Disk that stores numpy arrays as ``.npy`` bytes instead of pickles. + + :class:`SafeDisk` refuses numpy, because unpickling an array is not + memory-safe (see ``UNSAFE_MODULES``). :class:`SafeArrayDisk` keeps those + pickles out of the cache and stores arrays in numpy's own format: + + * An array on its own or inside a ``dict``, ``list`` or ``tuple`` is saved + with ``allow_pickle=False``, and its header is checked before it is + read back with ``allow_pickle=False``. + * numpy integers, floats, complex numbers, booleans and strings are stored + as the equivalent Python value, so they read back as that value. + * Arrays of objects, subclasses of ``ndarray`` such as masked arrays, and + arrays inside any other type raise :exc:`TypeError` when written. + + Everything else behaves as in :class:`SafeDisk`, including + :meth:`SafeDisk.with_types`. numpy stays optional: it is imported only to + read an array back. An encoded array is an ordinary pickled ``dict``, so + a process that uses another disk, or diskcache 5.6.3, reads that ``dict`` + instead of the array, and every process that shares a directory should + select this disk. + + """ + + def store(self, value, read, key=UNKNOWN): + numpy = sys.modules.get('numpy') + kinds = (numpy.ndarray, numpy.generic) if numpy else () + + def needs_encoding(item): + # pylint: disable=unidiomatic-typecheck + if type(item) is dict and _NPY_KEY in item: + raise TypeError( + 'SafeArrayDisk cannot store a dict with the key %r' + % (_NPY_KEY,) + ) + return isinstance(item, kinds) + + def encode(item): + # pylint: disable=unidiomatic-typecheck + if type(item) is numpy.ndarray: + return _pack_array(numpy, item) + if isinstance( + item, + (numpy.bool_, numpy.number, numpy.str_, numpy.bytes_), + ) and not isinstance(item, numpy.timedelta64): + return item.item() + return item + + if _contains(value, needs_encoding): + try: + value = _rebuild(value, encode) + except RecursionError: + raise TypeError( + 'SafeArrayDisk cannot store a nested value this deep' + ) from None + + return super().store(value, read, key=key) + + def fetch(self, mode, filename, value, read): + result = super().fetch(mode, filename, value, read) + + if not _contains(result, _encoded_array): + return result + + def decode(item): + if _encoded_array(item): + return _unpack_array(item[_NPY_KEY]) + return item + + try: + return _rebuild(result, decode) + except RecursionError: + raise pickle.UnpicklingError( + 'SafeArrayDisk refuses a value nested this deep' + ) from None diff --git a/diskcache_community/arrays.pyi b/diskcache_community/arrays.pyi new file mode 100644 index 0000000..ae596ce --- /dev/null +++ b/diskcache_community/arrays.pyi @@ -0,0 +1,3 @@ +from diskcache_community.core import SafeDisk + +class SafeArrayDisk(SafeDisk): ... diff --git a/diskcache_community/core.py b/diskcache_community/core.py index 5deec2d..ef11878 100644 --- a/diskcache_community/core.py +++ b/diskcache_community/core.py @@ -820,8 +820,9 @@ 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 + :class:`SafeArrayDisk ` to store + arrays. >>> import collections >>> Point = collections.namedtuple('Point', 'x y') diff --git a/docs/api.rst b/docs/api.rst index 85c419b..236682e 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -167,6 +167,16 @@ Read the :ref:`SafeJSONDisk tutorial ` for details. :special-members: :exclude-members: __weakref__ +SafeArrayDisk +------------- + +Read the :ref:`SafeArrayDisk tutorial ` for details. + +.. autoclass:: diskcache_community.SafeArrayDisk + :members: + :special-members: + :exclude-members: __weakref__ + Default disk ------------ diff --git a/docs/fix-cve-2025-69872.rst b/docs/fix-cve-2025-69872.rst index fc1a0b4..a662627 100644 --- a/docs/fix-cve-2025-69872.rst +++ b/docs/fix-cve-2025-69872.rst @@ -111,20 +111,16 @@ 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. Select + :class:`SafeArrayDisk ` to store arrays + instead, which saves them in numpy's ``.npy`` format: .. code-block:: python - import io + from diskcache_community import Cache, SafeArrayDisk - import numpy - - buffer = io.BytesIO() - numpy.save(buffer, array, allow_pickle=False) - cache['embedding'] = buffer.getvalue() - - array = numpy.load(io.BytesIO(cache['embedding']), allow_pickle=False) + cache = Cache('/var/tmp/mycache', disk=SafeArrayDisk) + cache['embedding'] = array 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..5fdcc23 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -1013,8 +1013,8 @@ 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 +crafted array pickle can corrupt memory. Use :class:`SafeArrayDisk +` 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. @@ -1056,6 +1056,48 @@ extend and :meth:`SafeDisk.with_types ``SafeJSONDisk.with_options(on_error='miss')`` to drop refused entries, as described in :ref:`tutorial-unreadable-entries`. +.. _tutorial-safearraydisk: + +SafeArrayDisk +............. + +:class:`SafeDisk ` refuses numpy, because +unpickling an array is not memory-safe. :class:`SafeArrayDisk +` stores arrays without pickling them: + +.. code-block:: python + + import numpy + + from diskcache_community import Cache, SafeArrayDisk + + cache = Cache('/tmp/mycache', disk=SafeArrayDisk) + cache['embedding'] = numpy.arange(3.0) + cache['result'] = {'scores': numpy.zeros(2), 'best': numpy.float64(1.5)} + +* An array on its own or inside a ``dict``, ``list`` or ``tuple`` is saved in + numpy's ``.npy`` format with ``allow_pickle=False``. Its header is checked + against the data before it is read back, also with ``allow_pickle=False``. +* numpy integers, floats, complex numbers, booleans and strings are stored as + the equivalent Python value, so ``numpy.float64(1.5)`` reads back as ``1.5``. +* Arrays of objects, subclasses of ``ndarray`` such as masked arrays and + ``numpy.matrix``, arrays inside any other type, and ``datetime64`` and + ``timedelta64`` scalars raise :exc:`TypeError` when written. + +Everything else behaves as in :class:`SafeDisk +`, and :meth:`SafeDisk.with_types +` and +``with_options(on_error='miss')`` work the same way. An entry whose array data +has been tampered with raises an error when read, or reads as missing with +``on_error='miss'``. + +The choice of disk is not stored in the cache directory, so select this disk in +every process that shares a directory. An encoded array is an ordinary pickled +``dict``: a process on another disk, or on diskcache 5.6.3, opens the directory +and reads that ``dict`` instead of the array. 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 diff --git a/tests/test_safearraydisk.py b/tests/test_safearraydisk.py new file mode 100644 index 0000000..7d2a6b6 --- /dev/null +++ b/tests/test_safearraydisk.py @@ -0,0 +1,543 @@ +"""Test diskcache_community.SafeArrayDisk.""" + +import collections as co +import io +import os +import os.path as op +import pickle +import shutil +import sqlite3 +import struct +import subprocess +import sys +import tempfile + +import pytest + +import diskcache_community as dc +from diskcache_community.arrays import _NPY_KEY +from diskcache_community.core import MODE_PICKLE + +np = pytest.importorskip('numpy') + +pytestmark = pytest.mark.filterwarnings('ignore', category=dc.EmptyDirWarning) + +MissDisk = dc.SafeArrayDisk.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,)) + + +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 write_row(directory, stored): + """Store `stored` in the value column of a pickle 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(pickle.dumps(stored, protocol=2))), + ) + con.commit() + con.close() + + +def npy(descr, shape, body=b'', version=1): + header = "{'descr': %s, 'fortran_order': False, 'shape': %s, }" % ( + descr, + shape, + ) + header = header.encode('latin1') + size = ' Date: Sun, 1 Nov 2026 21:18:48 -0800 Subject: [PATCH 02/10] Install numpy in the Django and SQLite CI jobs so the SafeArrayDisk tests run there. --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From 59e934f6d727a2e7330d18bcf0a17993502c50dd Mon Sep 17 00:00:00 2001 From: August Vi Date: Sun, 1 Nov 2026 21:59:31 -0800 Subject: [PATCH 03/10] Store SafeArrayDisk arrays through pickle persistent IDs and refuse forged entries with UnpicklingError. --- CHANGELOG.md | 3 +- CONTRIBUTING.md | 3 +- SECURITY.md | 3 +- diskcache_community/arrays.py | 325 ++++++++++-------- diskcache_community/core.py | 75 +++-- docs/fix-cve-2025-69872.rst | 3 + docs/tutorial.rst | 49 ++- tests/test_default_disk.py | 20 ++ tests/test_safearraydisk.py | 600 ++++++++++++++++++++++++++-------- 9 files changed, 755 insertions(+), 326 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 728e9c6..5950508 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,8 @@ 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). -- `SafeArrayDisk` is a `SafeDisk` that stores numpy arrays in the `.npy` format instead of pickling them, on their own or inside a `dict`, `list` or `tuple`, and stores numpy scalars as Python values. Array headers are checked against the data before reading, so a forged shape is refused rather than raising `MemoryError`. Object arrays and `ndarray` subclasses raise `TypeError` at write time. numpy stays optional and refused by `SafeDisk` (PR #34). +- `SafeArrayDisk` is a `SafeDisk` 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`. numpy scalars read back as Python values. 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 or by 5.6.3, which raise `UnpicklingError`. numpy stays optional, and `SafeDisk` still refuses it (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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e6c2134..2979e67 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -46,6 +46,7 @@ This fork publishes the same library under a new distribution name, `diskcache-c - 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 row format `SafeArrayDisk` writes: a persistent ID of `('npy', <.npy bytes>)` or `('item', )`, with `tests/test_safearraydisk.py` pinning it; - 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. +A process running 5.6.3 must be able to open a directory written by this fork, and the reverse. Entries a process wrote through `SafeArrayDisk` that hold numpy data are the one exception to reading them back: 5.6.3 opens the directory and raises `UnpicklingError` for those entries. Golden fixtures written by the 5.6.3 wheel guard this in the test suite. diff --git a/SECURITY.md b/SECURITY.md index 1303971..42ab948 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -19,9 +19,10 @@ 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. +- `SafeArrayDisk` is 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. +Unless you have opted into `SafeDisk`, `SafeJSONDisk` or `SafeArrayDisk`, 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. diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py index 3ff1843..d3a8918 100644 --- a/diskcache_community/arrays.py +++ b/diskcache_community/arrays.py @@ -5,105 +5,158 @@ import pickle import sys -from .core import UNKNOWN, SafeDisk +from .core import ( + SAFE_PICKLE_NAMES, + UNKNOWN, + Disk, + SafeDisk, + SafeUnpickler, + _ReaderMissing, +) -_NPY_KEY = '__diskcache_community_npy__' +# Pickle protocol 0 takes only strings as persistent IDs, and protocols below +# 3 have no opcode for bytes. +_MIN_PROTOCOL = 3 -def _contains(value, match): - """Return True if `match` holds for `value` or anything nested in it.""" - seen = set() - stack = [value] +def _check_header(numpy, data): + """Raise UnpicklingError unless the ``.npy`` bytes `data` are consistent. - while stack: - item = stack.pop() - if match(item): - return True - type_item = type(item) - if type_item in (dict, list, tuple): - if id(item) in seen: - continue - seen.add(id(item)) - if type_item is dict: - stack.extend(item.keys()) - stack.extend(item.values()) - else: - stack.extend(item) + The sizes in the header are checked against the data because a forged + shape would otherwise raise MemoryError, which ``on_error='miss'`` does + not treat as an unreadable entry. For the same reason, anything that goes + wrong while parsing the header, including a MemoryError from the Python + parser on a crafted one, is reported as UnpicklingError. - return False - - -def _rebuild(value, convert): - """Return `value` with `convert` applied through dicts, lists and tuples. + """ + # pylint: disable=broad-except,no-member + try: + reader = io.BytesIO(data) + version = numpy.lib.format.read_magic(reader) + if version == (1, 0): + header = numpy.lib.format.read_array_header_1_0(reader) + elif version == (2, 0): + header = numpy.lib.format.read_array_header_2_0(reader) + else: + raise ValueError('unsupported .npy version %r' % (version,)) + except Exception as error: + raise pickle.UnpicklingError( + 'SafeArrayDisk refuses the array header' + ) from error - `convert` returns its argument unchanged for anything it leaves alone. + shape, _, dtype = header + count = math.prod(shape) - """ - converted = convert(value) - if converted is not value: - return converted - type_value = type(value) - if type_value is dict: - return { - _rebuild(key, convert): _rebuild(item, convert) - for key, item in value.items() - } - if type_value is list: - return [_rebuild(item, convert) for item in value] - if type_value is tuple: - return tuple(_rebuild(item, convert) for item in value) - return value - - -def _encoded_array(item): - # pylint: disable=unidiomatic-typecheck - return type(item) is dict and len(item) == 1 and _NPY_KEY in item + # A zero-width dtype would let any shape match any amount of data. + if dtype.itemsize * count != len(data) - reader.tell() or ( + count and not dtype.itemsize + ): + raise pickle.UnpicklingError('SafeArrayDisk refuses the array header') def _pack_array(numpy, array): - """Return `array` as a dict that holds its ``.npy`` bytes.""" - if array.dtype.hasobject: - raise TypeError( - 'SafeArrayDisk cannot store an array of dtype %s' % (array.dtype,) - ) + """Return the ``.npy`` bytes of `array`, which a reader will accept.""" buffer = io.BytesIO() - numpy.save(buffer, array, allow_pickle=False) - data = buffer.getvalue() - if data[6:8] not in (b'\x01\x00', b'\x02\x00'): + try: + numpy.save(buffer, array, allow_pickle=False) + data = buffer.getvalue() + _check_header(numpy, data) + except (ValueError, pickle.UnpicklingError) as error: raise TypeError( 'SafeArrayDisk cannot store an array of dtype %s' % (array.dtype,) - ) - return {_NPY_KEY: data} + ) from error + return data def _unpack_array(data): - """Return the array in ``.npy`` bytes that may have been tampered with. - - The sizes in the header are checked against the data first, because a - forged shape would otherwise raise :exc:`MemoryError`, which - ``on_error='miss'`` does not treat as an unreadable entry. - - """ - # pylint: disable=import-error,import-outside-toplevel,no-member - import numpy - - reader = io.BytesIO(data) - version = numpy.lib.format.read_magic(reader) - if version == (1, 0): - header = numpy.lib.format.read_array_header_1_0(reader) - elif version == (2, 0): - header = numpy.lib.format.read_array_header_2_0(reader) - else: - raise pickle.UnpicklingError('SafeArrayDisk found no array data') - - shape, _, dtype = header - size = dtype.itemsize * math.prod(shape) - - if size != len(data) - reader.tell(): - raise pickle.UnpicklingError('SafeArrayDisk refuses the array header') - - reader.seek(0) - return numpy.lib.format.read_array(reader, allow_pickle=False) + """Return the array in ``.npy`` bytes that may have been tampered with.""" + # pylint: disable=broad-except,import-error,import-outside-toplevel + # pylint: disable=no-member + try: + import numpy + except ImportError as error: + raise _ReaderMissing( + 'SafeArrayDisk needs numpy to read an array' + ) from error + + _check_header(numpy, data) + + try: + return numpy.lib.format.read_array( + io.BytesIO(data), allow_pickle=False + ) + except MemoryError: + raise # The header matched the data, so this is a real shortage. + except Exception as error: + raise pickle.UnpicklingError( + 'SafeArrayDisk refuses the array data' + ) from error + + +def _item(numpy, scalar): + """Return the Python value of numpy `scalar`, or raise TypeError.""" + if isinstance( + scalar, (numpy.bool_, numpy.number, numpy.str_, numpy.bytes_) + ) and not isinstance(scalar, numpy.timedelta64): + value = scalar.item() + if not isinstance(value, numpy.generic): + return value + raise TypeError( + 'SafeArrayDisk cannot store a scalar of type numpy.%s' + % (type(scalar).__name__,) + ) + + +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 = {} + + 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( + 'SafeArrayDisk cannot store a subclass of ndarray: %s.%s' + % (type(obj).__module__, type(obj).__qualname__) + ) + pid = self._arrays.get(id(obj)) + if pid is None: + pid = ('npy', _pack_array(numpy, obj)) + self._arrays[id(obj)] = pid + return pid + if isinstance(obj, numpy.generic): + return ('item', _item(numpy, obj)) + return None + + +class _ArrayUnpickler(SafeUnpickler): + """Unpickler that reads back what _ArrayPickler wrote.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._arrays = {} + + def persistent_load(self, pid): + # pylint: disable=unidiomatic-typecheck + if type(pid) is tuple and len(pid) == 2: + kind, payload = pid + if kind == 'item': + return payload + if kind == 'npy': + # Decoding each bytes object once keeps what a row can decode + # to within its own size, and arrays that were shared shared. + # Holding `payload` keeps its id from being reused. + entry = self._arrays.get(id(payload)) + if entry is None: + entry = (payload, _unpack_array(payload)) + self._arrays[id(payload)] = entry + return entry[1] + raise pickle.UnpicklingError('SafeArrayDisk refuses a persistent ID') class SafeArrayDisk(SafeDisk): @@ -111,73 +164,65 @@ class SafeArrayDisk(SafeDisk): :class:`SafeDisk` refuses numpy, because unpickling an array is not memory-safe (see ``UNSAFE_MODULES``). :class:`SafeArrayDisk` keeps those - pickles out of the cache and stores arrays in numpy's own format: - - * An array on its own or inside a ``dict``, ``list`` or ``tuple`` is saved - with ``allow_pickle=False``, and its header is checked before it is - read back with ``allow_pickle=False``. - * numpy integers, floats, complex numbers, booleans and strings are stored - as the equivalent Python value, so they read back as that value. + pickles out of the cache by writing arrays through pickle's persistent IDs + in numpy's own format: + + * An array anywhere in a value is saved with ``allow_pickle=False``, and + its header is checked before it is read back with + ``allow_pickle=False``. That includes the containers :class:`SafeDisk` + accepts and the classes allowed with :meth:`SafeDisk.with_types`. + * numpy integers, floats, complex numbers, booleans and strings read back + as the equivalent Python value. * Arrays of objects, subclasses of ``ndarray`` such as masked arrays, and - arrays inside any other type raise :exc:`TypeError` when written. + other numpy scalars raise :exc:`TypeError` when written. + + Keys and everything else behave as in :class:`SafeDisk`. numpy stays + optional: it is used only for a value that needs it, and a process without + it raises :exc:`ImportError` for an entry with an array and keeps the + entry, even with ``on_error='miss'``. - Everything else behaves as in :class:`SafeDisk`, including - :meth:`SafeDisk.with_types`. numpy stays optional: it is imported only to - read an array back. An encoded array is an ordinary pickled ``dict``, so - a process that uses another disk, or diskcache 5.6.3, reads that ``dict`` - instead of the array, and every process that shares a directory should - select this disk. + 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. Every process that shares a directory should select this disk. """ + def _unpickler(self, data): + return _ArrayUnpickler( + io.BytesIO(data), + SAFE_PICKLE_NAMES, + self.safe_types, + self.safe_modules, + ) + def store(self, value, read, key=UNKNOWN): + # pylint: disable=unidiomatic-typecheck numpy = sys.modules.get('numpy') - kinds = (numpy.ndarray, numpy.generic) if numpy else () - def needs_encoding(item): - # pylint: disable=unidiomatic-typecheck - if type(item) is dict and _NPY_KEY in item: - raise TypeError( - 'SafeArrayDisk cannot store a dict with the key %r' - % (_NPY_KEY,) - ) - return isinstance(item, kinds) - - def encode(item): - # pylint: disable=unidiomatic-typecheck - if type(item) is numpy.ndarray: - return _pack_array(numpy, item) - if isinstance( - item, - (numpy.bool_, numpy.number, numpy.str_, numpy.bytes_), - ) and not isinstance(item, numpy.timedelta64): - return item.item() - return item - - if _contains(value, needs_encoding): - try: - value = _rebuild(value, encode) - except RecursionError: - raise TypeError( - 'SafeArrayDisk cannot store a nested value this deep' - ) from None - - return super().store(value, read, key=key) + if numpy is None or read or type(value) in (bytes, str, int, float): + return super().store(value, read, key=key) - def fetch(self, mode, filename, value, read): - result = super().fetch(mode, filename, value, read) + try: + self._check(value, 'value') + except TypeError: + pass + else: + return Disk.store(self, value, read, key=key) - if not _contains(result, _encoded_array): - return result + if isinstance(value, numpy.generic): + return super().store(_item(numpy, value), read, key=key) - def decode(item): - if _encoded_array(item): - return _unpack_array(item[_NPY_KEY]) - return item + buffer = io.BytesIO() + protocol = max(self.pickle_protocol, _MIN_PROTOCOL) try: - return _rebuild(result, decode) - except RecursionError: - raise pickle.UnpicklingError( - 'SafeArrayDisk refuses a value nested this deep' - ) from None + _ArrayPickler(buffer, protocol, numpy).dump(value) + except (pickle.PicklingError, AttributeError) as error: + raise TypeError( + 'SafeArrayDisk cannot store value: %s' % (error,) + ) from error + + data = buffer.getvalue() + self._check_pickle(data, value, 'value') + return self._store_pickle(data, value, key) diff --git a/diskcache_community/core.py b/diskcache_community/core.py index ef11878..79a9cca 100644 --- a/diskcache_community/core.py +++ b/diskcache_community/core.py @@ -347,13 +347,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) @@ -893,7 +896,12 @@ 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: @@ -1118,10 +1126,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 @@ -1172,27 +1180,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)): @@ -1501,9 +1508,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]' = [] diff --git a/docs/fix-cve-2025-69872.rst b/docs/fix-cve-2025-69872.rst index a662627..2dd18d4 100644 --- a/docs/fix-cve-2025-69872.rst +++ b/docs/fix-cve-2025-69872.rst @@ -122,6 +122,9 @@ Steps cache = Cache('/var/tmp/mycache', disk=SafeArrayDisk) cache['embedding'] = array + In the environment, set + ``DISKCACHE_DISK=diskcache_community:SafeArrayDisk``. The same lists apply. + 5. **Deal with entries already in the cache.** A cache written with the default :class:`Disk ` may hold objects that :class:`SafeDisk ` refuses, and reading them diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 5fdcc23..539121f 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -1063,7 +1063,8 @@ SafeArrayDisk :class:`SafeDisk ` refuses numpy, because unpickling an array is not memory-safe. :class:`SafeArrayDisk -` stores arrays without pickling them: +` stores arrays without pickling them, by +writing them through pickle's persistent IDs in numpy's ``.npy`` format: .. code-block:: python @@ -1075,26 +1076,38 @@ unpickling an array is not memory-safe. :class:`SafeArrayDisk cache['embedding'] = numpy.arange(3.0) cache['result'] = {'scores': numpy.zeros(2), 'best': numpy.float64(1.5)} -* An array on its own or inside a ``dict``, ``list`` or ``tuple`` is saved in - numpy's ``.npy`` format with ``allow_pickle=False``. Its header is checked - against the data before it is read back, also with ``allow_pickle=False``. -* numpy integers, floats, complex numbers, booleans and strings are stored as +* An array anywhere in a value is saved with ``allow_pickle=False``, including + inside classes allowed with :meth:`SafeDisk.with_types + `. Its header is checked against the + data before it is read back, also with ``allow_pickle=False``. An array that + appears twice in a value is stored once and is one array when read back. +* numpy integers, floats, complex numbers, booleans and strings read back as the equivalent Python value, so ``numpy.float64(1.5)`` reads back as ``1.5``. * Arrays of objects, subclasses of ``ndarray`` such as masked arrays and - ``numpy.matrix``, arrays inside any other type, and ``datetime64`` and - ``timedelta64`` scalars raise :exc:`TypeError` when written. + ``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 `. A memoized +function that returns a numpy scalar returns it unchanged on a miss and as a +Python value on a hit, so convert it in the function if a caller relies on its +type. Everything else behaves as in :class:`SafeDisk -`, and :meth:`SafeDisk.with_types -` and -``with_options(on_error='miss')`` work the same way. An entry whose array data -has been tampered with raises an error when read, or reads as missing with +`, 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'``. -The choice of disk is not stored in the cache directory, so select this disk in -every process that shares a directory. An encoded array is an ordinary pickled -``dict``: a process on another disk, or on diskcache 5.6.3, opens the directory -and reads that ``dict`` instead of the array. Arrays that the default +A value with no numpy in it is an ordinary pickle. An entry holding numpy data +is not readable by any other disk, 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. +A process that cannot import numpy raises :exc:`ImportError` for an entry with +an array 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. @@ -1123,8 +1136,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 +``diskcache_community:SafeArrayDisk``, ``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 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_safearraydisk.py b/tests/test_safearraydisk.py index 7d2a6b6..f5664b4 100644 --- a/tests/test_safearraydisk.py +++ b/tests/test_safearraydisk.py @@ -1,10 +1,11 @@ """Test diskcache_community.SafeArrayDisk.""" import collections as co +import dataclasses import io -import os import os.path as op import pickle +import pickletools import shutil import sqlite3 import struct @@ -15,7 +16,7 @@ import pytest import diskcache_community as dc -from diskcache_community.arrays import _NPY_KEY +from diskcache_community import arrays from diskcache_community.core import MODE_PICKLE np = pytest.importorskip('numpy') @@ -42,6 +43,24 @@ def __reduce__(self): return (_execute, (self.path,)) +class Outside: + """A class outside the allowlist.""" + + +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') @@ -59,24 +78,47 @@ def same(actual, expected): assert actual.tobytes() == expected.tobytes() -def write_row(directory, stored): - """Store `stored` in the value column of a pickle row.""" +def npy_bytes(array): + buffer = io.BytesIO() + np.save(buffer, array, allow_pickle=False) + return buffer.getvalue() + + +def persid(pid): + """Return a pickle of one persistent ID, without using the code tested.""" + return pickle.dumps(pid, protocol=3)[:-1] + pickle.BINPERSID + pickle.STOP + + +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(pickle.dumps(stored, protocol=2))), + (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=dc.SafeArrayDisk, 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', 'expected'), [ @@ -158,6 +296,12 @@ def test_scalars_become_python_values(directory, scalar, expected): assert type(cache['nested']['metrics'][0]) is type(expected) +def test_a_top_level_integer_scalar_is_a_plain_row(directory): + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + cache['key'] = np.int64(2) + assert cache.incr('key') == 3 + + def test_scalar_dict_keys_become_python_values(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: cache['key'] = {np.int64(1): 'one'} @@ -165,6 +309,31 @@ def test_scalar_dict_keys_become_python_values(directory): assert type(next(iter(cache['key']))) is int +@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=dc.SafeArrayDisk) 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 int + + @pytest.mark.parametrize( 'scalar', [ @@ -177,6 +346,19 @@ def test_time_scalars_are_refused(directory, scalar): with dc.Cache(directory, disk=dc.SafeArrayDisk) 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=dc.SafeArrayDisk) as cache: + try: + cache['key'] = {'a': value} + except TypeError: + return + assert cache['key'] == {'a': 1.5} @pytest.mark.parametrize( @@ -195,12 +377,28 @@ def test_object_arrays_are_refused(directory, array): assert len(cache) == 0 -def test_field_names_that_need_format_3_are_refused(directory): - array = np.zeros(1, dtype=[('α', 'i4')]) - - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: +@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( @@ -215,87 +413,168 @@ def test_field_names_that_need_format_3_are_refused(directory): def test_subclasses_are_refused(directory, make): """numpy.save would silently drop what makes them a subclass.""" with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: - with pytest.raises(TypeError): + with pytest.raises(TypeError, match='subclass'): cache['key'] = make() + with pytest.raises(TypeError, match='subclass'): + cache['key'] = [make()] -@pytest.mark.parametrize( - 'make', - [ - lambda array: Pair(array, 1), - lambda array: co.OrderedDict(a=array), - lambda array: co.deque([array]), - ], - ids=['namedtuple', 'ordereddict', 'deque'], -) -def test_arrays_in_other_containers_are_refused(directory, make): +def test_numpy_in_a_key_is_refused(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: with pytest.raises(TypeError): - cache['key'] = make(np.arange(2.0)) + 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=dc.SafeArrayDisk) 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_is_numpy_on_a_miss_and_python_on_a_hit(directory): + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + + @cache.memoize() + def compute(): + return np.int64(3) + + assert type(compute()) is np.int64 + assert type(compute()) is int -def test_array_key_is_refused(directory): +def test_memoize_scalar_argument_is_refused(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + + @cache.memoize() + def compute(value): + return value + with pytest.raises(TypeError): - cache[(np.int64(1),)] = 'value' + compute(np.int64(3)) -def test_reserved_key_is_refused(directory): +def test_plain_values_behave_as_in_safedisk(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) 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['key'] = {_NPY_KEY: b'data'} + cache['outside'] = Outside() with pytest.raises(TypeError): - cache['key'] = [{'a': {_NPY_KEY: 1, 'other': 2}}] + cache['function'] = lambda: None + +def test_plain_values_are_ordinary_pickles(directory): + with dc.Cache(directory, disk=dc.SafeArrayDisk) 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_reserved_key_is_refused_without_numpy(directory): + +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=dc.SafeArrayDisk) 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' - 'from diskcache_community.arrays import _NPY_KEY\n' 'assert "numpy" not in sys.modules\n' 'with dc.Cache(sys.argv[1], disk=dc.SafeArrayDisk) as cache:\n' ' cache["plain"] = {"a": [1, 2]}\n' ' assert cache["plain"] == {"a": [1, 2]}\n' - ' try:\n' - ' cache["key"] = {_NPY_KEY: b"data"}\n' - ' except TypeError:\n' - ' pass\n' - ' else:\n' - ' raise AssertionError("not refused")\n' 'assert "numpy" not in sys.modules\n' ) subprocess.run([sys.executable, '-c', code, directory], check=True) -def test_plain_values_behave_as_in_safedisk(directory): +def test_reads_an_array_in_a_process_that_has_not_imported_numpy(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) 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['object'] = object() + cache['key'] = np.arange(2.0) + code = ( + 'import sys\n' + 'import diskcache_community as dc\n' + 'assert "numpy" not in sys.modules\n' + 'with dc.Cache(sys.argv[1], disk=dc.SafeArrayDisk) as cache:\n' + ' assert cache["key"].tolist() == [0.0, 1.0]\n' + ) + subprocess.run([sys.executable, '-c', code, directory], check=True) -def test_cyclic_values_still_work(directory): - value = [1] - value.append(value) +def test_without_numpy_an_entry_raises_and_is_kept(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: - cache['key'] = value - result = cache['key'] + cache['key'] = np.arange(2.0) - assert result[0] == 1 - assert result[1] is result + code = ( + 'import sys\n' + 'sys.modules["numpy"] = None\n' + 'import diskcache_community as dc\n' + 'disk = dc.SafeArrayDisk.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=dc.SafeArrayDisk) as cache: + same(cache['key'], np.arange(2.0)) -def test_cycle_with_an_array_is_refused(directory): - value = [np.arange(2.0)] - value.append(value) +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)) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: - with pytest.raises(TypeError): - cache['key'] = value + +def test_real_memory_shortage_keeps_the_entry(directory, monkeypatch): + def fail(*args, **kwargs): + raise MemoryError + + with dc.Cache(directory, disk=MissDisk) as cache: + cache['key'] = np.arange(2.0) + monkeypatch.setattr(np.lib.format, 'read_array', fail) + with pytest.raises(MemoryError): + cache.get('key') + monkeypatch.undo() + same(cache['key'], np.arange(2.0)) def test_read_true_still_returns_a_file(directory): @@ -316,22 +595,6 @@ def test_large_array_lives_in_a_file(directory): assert cache.check() == [] -def test_memoize(directory): - calls = [] - - with dc.Cache(directory, disk=dc.SafeArrayDisk) 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_fanout_index_and_deque(directory): array = np.arange(3.0) @@ -373,13 +636,57 @@ def test_with_types_still_refuses_numpy(): dc.SafeArrayDisk.with_types(np.ndarray) +def test_row_format_is_pinned(directory): + """Changing what is on disk breaks every cache: this is hand-built.""" + array = np.arange(3.0) + write_row(directory, persid(('npy', npy_bytes(array)))) + + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + same(cache['key'], array) + + other = op.join(directory, 'other') + write_row(other, persid(('item', 1.5))) + + with dc.Cache(other, disk=dc.SafeArrayDisk) as cache: + assert cache['key'] == 1.5 + + +def test_stored_row_uses_the_pinned_format(directory): + array = np.arange(3.0) + + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + cache['array'] = array + cache['scalar'] = [np.int64(1)] + rows = row_bytes(cache, 'array'), row_bytes(cache, 'scalar') + + ops = [(op.name, arg) for op, arg, _ in pickletools.genops(rows[0])] + assert ('BINPERSID', None) in ops + assert any(arg == 'npy' for _, arg in ops) + assert any(arg == npy_bytes(array) for _, arg in ops) + assert any(arg == 'item' for _, arg, _ in pickletools.genops(rows[1])) + + @pytest.mark.parametrize('reader', [dc.Disk, dc.SafeDisk]) -def test_another_disk_reads_a_dict(directory, reader): +def test_another_disk_refuses_an_array_entry(directory, reader): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: - cache['key'] = np.arange(2.0) + cache['array'] = np.arange(2.0) + cache['plain'] = {'a': 1} with dc.Cache(directory, disk=reader) as cache: - assert list(cache['key']) == [_NPY_KEY] + assert cache['plain'] == {'a': 1} + with pytest.raises(pickle.UnpicklingError): + cache['array'] + + +def test_another_disk_with_miss_deletes_an_array_entry(directory): + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + cache['array'] = np.arange(2.0) + + disk = dc.SafeDisk.with_options(on_error='miss') + + with dc.Cache(directory, disk=disk) as cache: + assert cache.get('array', 'default') == 'default' + assert len(cache) == 0 def test_reads_safedisk_entries(directory): @@ -404,7 +711,17 @@ def test_array_pickled_by_disk_is_refused(directory): def test_poison_does_not_run(directory): marker = directory + '.executed' - write_row(directory, Poison(marker)) + write_row(directory, pickle.dumps(Poison(marker), protocol=2)) + + with dc.Cache(directory, disk=MissDisk) as cache: + assert cache.get('key', 'default') == 'default' + + assert not op.exists(marker) + + +def test_poison_inside_a_persistent_id_does_not_run(directory): + marker = directory + '.executed' + write_row(directory, persid(('item', Poison(marker)))) with dc.Cache(directory, disk=MissDisk) as cache: assert cache.get('key', 'default') == 'default' @@ -412,9 +729,9 @@ def test_poison_does_not_run(directory): assert not op.exists(marker) -def test_poison_inside_the_array_slot_does_not_run(directory): +def test_pickle_in_the_array_slot_does_not_run(directory): marker = directory + '.executed' - write_row(directory, {_NPY_KEY: pickle.dumps(Poison(marker))}) + write_row(directory, persid(('npy', pickle.dumps(Poison(marker))))) with dc.Cache(directory, disk=MissDisk) as cache: assert cache.get('key', 'default') == 'default' @@ -428,44 +745,44 @@ def test_poison_inside_the_array_slot_does_not_run(directory): 'pickle': pickle.dumps([1, 2]), 'zip magic': b'PK\x03\x04' + b'\x00' * 40, 'npz-like': b'PK\x05\x06' + b'\x00' * 18, - 'str': 'text', - 'int': 1, - 'none': None, - 'list': [1, 2], 'truncated': npy("' Date: Sun, 1 Nov 2026 22:17:06 -0800 Subject: [PATCH 04/10] Keep numpy NaN scalars in SafeArrayDisk instead of letting SQLite store them as NULL. --- diskcache_community/arrays.py | 5 ++++- docs/tutorial.rst | 4 +++- tests/test_safearraydisk.py | 12 ++++++++++++ 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py index d3a8918..f258d7e 100644 --- a/diskcache_community/arrays.py +++ b/diskcache_community/arrays.py @@ -211,7 +211,10 @@ def store(self, value, read, key=UNKNOWN): return Disk.store(self, value, read, key=key) if isinstance(value, numpy.generic): - return super().store(_item(numpy, value), read, key=key) + item = _item(numpy, value) + # SQLite stores a raw NaN as NULL, so NaN stays in a pickle below. + if item == item: # pylint: disable=comparison-with-itself + return super().store(item, read, key=key) buffer = io.BytesIO() protocol = max(self.pickle_protocol, _MIN_PROTOCOL) diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 539121f..8a15cfb 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -1169,7 +1169,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_safearraydisk.py b/tests/test_safearraydisk.py index f5664b4..559300a 100644 --- a/tests/test_safearraydisk.py +++ b/tests/test_safearraydisk.py @@ -302,6 +302,18 @@ def test_a_top_level_integer_scalar_is_a_plain_row(directory): assert cache.incr('key') == 3 +@pytest.mark.parametrize('kind', [np.float64, np.float32, np.float16]) +def test_nan_scalar_is_not_lost(directory, kind): + """A raw float NaN would come back as None, because SQLite stores NULL.""" + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + cache['key'] = kind('nan') + cache['inf'] = kind('inf') + result = cache['key'] + assert type(result) is float + assert result != result + assert cache['inf'] == float('inf') + + def test_scalar_dict_keys_become_python_values(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: cache['key'] = {np.int64(1): 'one'} From c94e2dc4caef07b07cfbcd9eed119a02410aeb64 Mon Sep 17 00:00:00 2001 From: August Vi Date: Sun, 1 Nov 2026 22:37:27 -0800 Subject: [PATCH 05/10] Hold arrays in the SafeArrayDisk write memo and keep persistent IDs out of key rows. --- diskcache_community/arrays.py | 18 ++++++++---- tests/test_safearraydisk.py | 52 +++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py index f258d7e..6d04928 100644 --- a/diskcache_community/arrays.py +++ b/diskcache_community/arrays.py @@ -124,11 +124,13 @@ def persistent_id(self, obj): 'SafeArrayDisk cannot store a subclass of ndarray: %s.%s' % (type(obj).__module__, type(obj).__qualname__) ) - pid = self._arrays.get(id(obj)) - if pid is None: - pid = ('npy', _pack_array(numpy, obj)) - self._arrays[id(obj)] = pid - return pid + # Holding `obj` keeps its id from being reused: pickle's own memo + # skips persistent IDs, and a value may hand it temporary arrays. + entry = self._arrays.get(id(obj)) + if entry is None: + entry = (obj, ('npy', _pack_array(numpy, obj))) + self._arrays[id(obj)] = entry + return entry[1] if isinstance(obj, numpy.generic): return ('item', _item(numpy, obj)) return None @@ -196,6 +198,12 @@ def _unpickler(self, data): self.safe_modules, ) + def get(self, key, raw): + if raw: + return super().get(key, raw) + # Keys never hold numpy, so a key row gets no persistent IDs. + return super()._unpickler(key).load() + def store(self, value, read, key=UNKNOWN): # pylint: disable=unidiomatic-typecheck numpy = sys.modules.get('numpy') diff --git a/tests/test_safearraydisk.py b/tests/test_safearraydisk.py index 559300a..23d2c68 100644 --- a/tests/test_safearraydisk.py +++ b/tests/test_safearraydisk.py @@ -47,6 +47,17 @@ 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.""" @@ -234,6 +245,16 @@ def test_shared_arrays_stay_shared_and_are_stored_once(directory): assert extra < 1000 +def test_temporary_arrays_keep_their_own_data(directory): + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + cache['key'] = Temporaries(50) + result = cache['key'] + + assert [item.tolist() for item in result] == [ + [n, n + 1.0] for n in range(50) + ] + + def test_cycles_work(directory): array = np.arange(2.0) value = [array] @@ -832,6 +853,37 @@ def test_forged_persistent_id_is_refused(directory, pid): assert len(cache) == 0 +@pytest.mark.parametrize( + 'pid', + [('npy', npy_bytes(np.arange(2.0))), ('item', 1)], + ids=['array', 'item'], +) +def test_persistent_id_in_a_key_row_is_refused(directory, pid): + 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)),), + ) + con.commit() + con.close() + + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with pytest.raises(pickle.UnpicklingError): + list(cache) + + with dc.Cache(directory, disk=MissDisk) as cache: + assert list(cache) == [] + + +def test_keys_still_read_back(directory): + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + cache[('a', 1)] = np.arange(2.0) + cache['text'] = np.arange(2.0) + assert list(cache) == [('a', 1), 'text'] + + def test_forged_array_nested_in_a_value_is_a_miss(directory): inner = persid(('npy', FORGED['huge shape'])) write_row(directory, b'\x80\x03]' + inner[2:-1] + b'a.') From 694119140a9fcd4978a6688f1ce531b24dc8d6b4 Mon Sep 17 00:00:00 2001 From: August Vi Date: Sun, 1 Nov 2026 23:09:22 -0800 Subject: [PATCH 06/10] Keep numpy scalar types, read large arrays without a second copy and keep float NaN on every disk. --- CHANGELOG.md | 3 +- CONTRIBUTING.md | 2 +- diskcache_community/arrays.py | 228 +++++++++++++++++++---------- diskcache_community/core.py | 21 +-- docs/tutorial.rst | 17 +-- tests/test_core.py | 25 ++++ tests/test_safearraydisk.py | 268 ++++++++++++++++++++++++---------- tests/test_safedisk.py | 15 ++ 8 files changed, 406 insertions(+), 173 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5950508..b9462ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,11 +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). -- `SafeArrayDisk` is a `SafeDisk` 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`. numpy scalars read back as Python values. 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 or by 5.6.3, which raise `UnpicklingError`. numpy stays optional, and `SafeDisk` still refuses it (PR #34). +- `SafeArrayDisk` is a `SafeDisk` 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`. 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 or by 5.6.3, which raise `UnpicklingError`. numpy stays optional, and `SafeDisk` still refuses it (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 2979e67..3a3debe 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -46,7 +46,7 @@ This fork publishes the same library under a new distribution name, `diskcache-c - 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 row format `SafeArrayDisk` writes: a persistent ID of `('npy', <.npy bytes>)` or `('item', )`, with `tests/test_safearraydisk.py` pinning it; +- the row format `SafeArrayDisk` writes: a persistent ID of `('npy', )` or `('scalar', , )` in a pickle of protocol 5 or higher, with `tests/test_safearraydisk.py` pinning it; - 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. Entries a process wrote through `SafeArrayDisk` that hold numpy data are the one exception to reading them back: 5.6.3 opens the directory and raises `UnpicklingError` for those entries. Golden fixtures written by the 5.6.3 wheel guard this in the test suite. diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py index 6d04928..4ed8709 100644 --- a/diskcache_community/arrays.py +++ b/diskcache_community/arrays.py @@ -14,24 +14,61 @@ _ReaderMissing, ) -# Pickle protocol 0 takes only strings as persistent IDs, and protocols below -# 3 have no opcode for bytes. -_MIN_PROTOCOL = 3 +# Pickle protocol 5 is the first that stores a bytearray in place, so that +# reading a file-backed entry needs no second copy of its array. +_MIN_PROTOCOL = 5 + +# Bytes of a ``.npy`` file that can hold its header: the prefix, plus the +# longest header that numpy parses by default, plus padding. +_HEADER_LIMIT = 12 + 10000 + 64 + +# The numpy scalar types that are stored, by the name written to the row. The +# names do not vary by platform or numpy version, unlike the types' own names. +_SCALARS = { + 'bool': 'bool_', + 'int8': 'int8', + 'int16': 'int16', + 'int32': 'int32', + 'int64': 'int64', + 'uint8': 'uint8', + 'uint16': 'uint16', + 'uint32': 'uint32', + 'uint64': 'uint64', + 'float16': 'float16', + 'float32': 'float32', + 'float64': 'float64', + 'complex64': 'complex64', + 'complex128': 'complex128', + 'str': 'str_', + 'bytes': 'bytes_', +} +_SCALAR_VALUES = (bool, int, float, complex, str, bytes) + + +def _import_numpy(): + # pylint: disable=import-error,import-outside-toplevel + try: + import numpy + except ImportError as error: + raise _ReaderMissing( + 'SafeArrayDisk needs numpy to read this entry' + ) from error + return numpy -def _check_header(numpy, data): - """Raise UnpicklingError unless the ``.npy`` bytes `data` are consistent. +def _read_header(numpy, data): + """Return shape, Fortran order, dtype and offset of ``.npy`` bytes `data`. - The sizes in the header are checked against the data because a forged - shape would otherwise raise MemoryError, which ``on_error='miss'`` does - not treat as an unreadable entry. For the same reason, anything that goes - wrong while parsing the header, including a MemoryError from the Python - parser on a crafted one, is reported as UnpicklingError. + Raise UnpicklingError unless the header is consistent with the 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(data) + reader = io.BytesIO(bytes(memoryview(data)[:_HEADER_LIMIT])) version = numpy.lib.format.read_magic(reader) if version == (1, 0): header = numpy.lib.format.read_array_header_1_0(reader) @@ -44,67 +81,87 @@ def _check_header(numpy, data): 'SafeArrayDisk refuses the array header' ) from error - shape, _, dtype = header + shape, fortran, dtype = header + offset = reader.tell() count = math.prod(shape) # A zero-width dtype would let any shape match any amount of data. - if dtype.itemsize * count != len(data) - reader.tell() or ( - count and not dtype.itemsize + if ( + dtype.hasobject + or dtype.itemsize * count != len(data) - offset + or (count and not dtype.itemsize) ): raise pickle.UnpicklingError('SafeArrayDisk refuses the array header') + return shape, fortran, dtype, offset + def _pack_array(numpy, array): - """Return the ``.npy`` bytes of `array`, which a reader will accept.""" + """Return a buffer of the ``.npy`` bytes of `array`, for pickle to store. + + The buffer is writable, so pickle stores it as a bytearray in place. + + """ buffer = io.BytesIO() try: numpy.save(buffer, array, allow_pickle=False) - data = buffer.getvalue() - _check_header(numpy, data) + view = buffer.getbuffer() + _read_header(numpy, view) except (ValueError, pickle.UnpicklingError) as error: raise TypeError( 'SafeArrayDisk cannot store an array of dtype %s' % (array.dtype,) ) from error - return data + return pickle.PickleBuffer(view) def _unpack_array(data): - """Return the array in ``.npy`` bytes that may have been tampered with.""" - # pylint: disable=broad-except,import-error,import-outside-toplevel - # pylint: disable=no-member - try: - import numpy - except ImportError as error: - raise _ReaderMissing( - 'SafeArrayDisk needs numpy to read an array' - ) from error + """Return the array in the bytearray `data`, which it keeps as its data. - _check_header(numpy, data) + A view takes no second copy of the array, and stays writable because + `data` is. `data` may have been tampered with. + + """ + # pylint: disable=broad-except + numpy = _import_numpy() + shape, fortran, dtype, offset = _read_header(numpy, data) try: - return numpy.lib.format.read_array( - io.BytesIO(data), allow_pickle=False - ) - except MemoryError: - raise # The header matched the data, so this is a real shortage. + if not math.prod(shape): + return numpy.empty(shape, dtype) + array = numpy.frombuffer(data, dtype, offset=offset) + return array.reshape(shape, order='F' if fortran else 'C') except Exception as error: raise pickle.UnpicklingError( 'SafeArrayDisk refuses the array data' ) from error -def _item(numpy, scalar): - """Return the Python value of numpy `scalar`, or raise TypeError.""" - if isinstance( - scalar, (numpy.bool_, numpy.number, numpy.str_, numpy.bytes_) - ) and not isinstance(scalar, numpy.timedelta64): - value = scalar.item() - if not isinstance(value, numpy.generic): - return value - raise TypeError( - 'SafeArrayDisk cannot store a scalar of type numpy.%s' - % (type(scalar).__name__,) - ) +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( + 'SafeArrayDisk 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.""" + # pylint: disable=broad-except,unidiomatic-typecheck + if type(value) not in _SCALAR_VALUES: + raise pickle.UnpicklingError('SafeArrayDisk refuses a scalar') + + numpy = _import_numpy() + + try: + return getattr(numpy, _SCALARS[name])(value) + except Exception as error: + raise pickle.UnpicklingError( + 'SafeArrayDisk refuses a scalar' + ) from error class _ArrayPickler(pickle.Pickler): @@ -114,6 +171,14 @@ 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.""" + # pylint: disable=unused-argument + self.reduced = True + return NotImplemented def persistent_id(self, obj): # pylint: disable=unidiomatic-typecheck @@ -132,7 +197,17 @@ def persistent_id(self, obj): self._arrays[id(obj)] = entry return entry[1] if isinstance(obj, numpy.generic): - return ('item', _item(numpy, obj)) + 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( + 'SafeArrayDisk cannot store a scalar of type numpy.%s' + % (kind.__name__,) + ) + return ('scalar', name, value) return None @@ -145,19 +220,20 @@ def __init__(self, *args, **kwargs): def persistent_load(self, pid): # pylint: disable=unidiomatic-typecheck - if type(pid) is tuple and len(pid) == 2: - kind, payload = pid - if kind == 'item': - return payload - if kind == 'npy': - # Decoding each bytes object once keeps what a row can decode - # to within its own size, and arrays that were shared shared. - # Holding `payload` keeps its id from being reused. - entry = self._arrays.get(id(payload)) - if entry is None: - entry = (payload, _unpack_array(payload)) - self._arrays[id(payload)] = entry - return entry[1] + if type(pid) is tuple: + if len(pid) == 3 and pid[0] == 'scalar': + return _unpack_scalar(pid[1], pid[2]) + if len(pid) == 2 and pid[0] == 'npy': + payload = pid[1] + if type(payload) is bytearray: + # Decoding each buffer once keeps what a row can decode to + # within its own size, and shared arrays shared. Holding + # `payload` keeps its id from being reused. + entry = self._arrays.get(id(payload)) + if entry is None: + entry = (payload, _unpack_array(payload)) + self._arrays[id(payload)] = entry + return entry[1] raise pickle.UnpicklingError('SafeArrayDisk refuses a persistent ID') @@ -170,17 +246,17 @@ class SafeArrayDisk(SafeDisk): in numpy's own format: * An array anywhere in a value is saved with ``allow_pickle=False``, and - its header is checked before it is read back with - ``allow_pickle=False``. That includes the containers :class:`SafeDisk` - accepts and the classes allowed with :meth:`SafeDisk.with_types`. - * numpy integers, floats, complex numbers, booleans and strings read back - as the equivalent Python value. + its header is checked before it is read back. That includes the + containers :class:`SafeDisk` accepts and the classes allowed with + :meth:`SafeDisk.with_types`. + * numpy booleans, integers, floats, complex numbers and strings read back + as the same numpy type. * Arrays of objects, subclasses of ``ndarray`` such as masked arrays, and other numpy scalars raise :exc:`TypeError` when written. Keys and everything else behave as in :class:`SafeDisk`. numpy stays optional: it is used only for a value that needs it, and a process without - it raises :exc:`ImportError` for an entry with an array and keeps the + it raises :exc:`ImportError` for an entry holding numpy data and keeps the entry, even with ``on_error='miss'``. A value with no numpy in it is an ordinary pickle. An entry holding numpy @@ -190,9 +266,9 @@ class SafeArrayDisk(SafeDisk): """ - def _unpickler(self, data): + def _unpickler(self, file): return _ArrayUnpickler( - io.BytesIO(data), + file, SAFE_PICKLE_NAMES, self.safe_types, self.safe_modules, @@ -202,7 +278,7 @@ def get(self, key, raw): if raw: return super().get(key, raw) # Keys never hold numpy, so a key row gets no persistent IDs. - return super()._unpickler(key).load() + return super()._unpickler(io.BytesIO(key)).load() def store(self, value, read, key=UNKNOWN): # pylint: disable=unidiomatic-typecheck @@ -218,22 +294,22 @@ def store(self, value, read, key=UNKNOWN): else: return Disk.store(self, value, read, key=key) - if isinstance(value, numpy.generic): - item = _item(numpy, value) - # SQLite stores a raw NaN as NULL, so NaN stays in a pickle below. - if item == item: # pylint: disable=comparison-with-itself - return super().store(item, read, key=key) - buffer = io.BytesIO() protocol = max(self.pickle_protocol, _MIN_PROTOCOL) + pickler = _ArrayPickler(buffer, protocol, numpy) try: - _ArrayPickler(buffer, protocol, numpy).dump(value) + pickler.dump(value) except (pickle.PicklingError, AttributeError) as error: raise TypeError( 'SafeArrayDisk cannot store value: %s' % (error,) ) from error data = buffer.getvalue() - self._check_pickle(data, value, 'value') + + # Only a reduced object can look up a name, and scanning the pickle + # would copy every array in it. + if pickler.reduced: + self._check_pickle(data, value, 'value') + return self._store_pickle(data, value, key) diff --git a/diskcache_community/core.py b/diskcache_community/core.py index 79a9cca..df5cdd4 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: @@ -875,9 +877,9 @@ def with_types(cls, *types, modules=()): }, ) - def _unpickler(self, data): + def _unpickler(self, file): return SafeUnpickler( - io.BytesIO(data), + file, SAFE_PICKLE_NAMES, self.safe_types, self.safe_modules, @@ -906,7 +908,7 @@ def _check_pickle(self, data, obj, what): 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) @@ -915,7 +917,7 @@ def _check_pickle(self, data, 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' @@ -1012,7 +1014,7 @@ def put(self, key): def get(self, key, raw): if raw: return super().get(key, raw) - return self._unpickler(key).load() + return self._unpickler(io.BytesIO(key)).load() def store(self, value, read, key=UNKNOWN): # pylint: disable=unidiomatic-typecheck @@ -1026,8 +1028,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) diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 8a15cfb..b92a96a 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -1081,8 +1081,8 @@ writing them through pickle's persistent IDs in numpy's ``.npy`` format: `. Its header is checked against the data before it is read back, also with ``allow_pickle=False``. An array that appears twice in a value is stored once and is one array when read back. -* numpy integers, floats, complex numbers, booleans and strings read back as - the equivalent Python value, so ``numpy.float64(1.5)`` reads back as ``1.5``. +* 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 @@ -1090,10 +1090,7 @@ writing them through pickle's persistent IDs in numpy's ``.npy`` format: 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 `. A memoized -function that returns a numpy scalar returns it unchanged on a miss and as a -Python value on a hit, so convert it in the function if a caller relies on its -type. +decorated with :meth:`memoize `. Everything else behaves as in :class:`SafeDisk `, and ``with_options(on_error='miss')`` works @@ -1106,10 +1103,10 @@ is not readable by any other disk, 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. -A process that cannot import numpy raises :exc:`ImportError` for an entry with -an array 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. +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: 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_safearraydisk.py b/tests/test_safearraydisk.py index 23d2c68..dfc00a6 100644 --- a/tests/test_safearraydisk.py +++ b/tests/test_safearraydisk.py @@ -12,6 +12,7 @@ import subprocess import sys import tempfile +import tracemalloc import pytest @@ -95,9 +96,32 @@ def npy_bytes(array): return buffer.getvalue() +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 a pickle of one persistent ID, without using the code tested.""" - return pickle.dumps(pid, protocol=3)[:-1] + pickle.BINPERSID + pickle.STOP + return forge(Pid(pid)) + + +def array_pid(data): + return ('npy', bytearray(data)) def write_row(directory, data): @@ -295,51 +319,53 @@ def test_every_pickle_protocol_setting(directory, protocol): @pytest.mark.parametrize( - ('scalar', 'expected'), + 'scalar', [ - (np.float64(1.5), 1.5), - (np.float32(1.5), 1.5), - (np.int64(2), 2), - (np.uint8(2), 2), - (np.bool_(True), True), - (np.complex128(1j), 1j), - (np.str_('text'), 'text'), - (np.bytes_(b'data'), b'data'), + 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_become_python_values(directory, scalar, expected): +def test_scalars_keep_their_type_and_value(directory, scalar): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: cache['key'] = scalar - cache['nested'] = {'metrics': [scalar]} - assert cache['key'] == expected - assert type(cache['key']) is type(expected) - assert type(cache['nested']['metrics'][0]) is type(expected) - - -def test_a_top_level_integer_scalar_is_a_plain_row(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: - cache['key'] = np.int64(2) - assert cache.incr('key') == 3 + 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) @pytest.mark.parametrize('kind', [np.float64, np.float32, np.float16]) -def test_nan_scalar_is_not_lost(directory, kind): - """A raw float NaN would come back as None, because SQLite stores NULL.""" +def test_scalar_nan_and_infinity_are_kept(directory, kind): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: - cache['key'] = kind('nan') + cache['nan'] = kind('nan') cache['inf'] = kind('inf') - result = cache['key'] - assert type(result) is float - assert result != result - assert cache['inf'] == float('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_scalar_dict_keys_become_python_values(directory): +def test_scalars_as_dict_keys_keep_their_type(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: cache['key'] = {np.int64(1): 'one'} assert cache['key'] == {1: 'one'} - assert type(next(iter(cache['key']))) is int + assert type(next(iter(cache['key']))) is np.int64 @pytest.mark.parametrize( @@ -364,7 +390,7 @@ def test_scalars_in_the_other_safe_containers(directory, make): assert type(result) is type(value) item = next(iter(result.values() if hasattr(result, 'values') else result)) assert item == 2 - assert type(item) is int + assert type(item) is np.int64 @pytest.mark.parametrize( @@ -476,15 +502,17 @@ def compute(size): assert calls == [3] -def test_memoize_scalar_is_numpy_on_a_miss_and_python_on_a_hit(directory): +def test_memoize_scalar_keeps_its_type_on_a_hit(directory): with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: @cache.memoize() def compute(): - return np.int64(3) + return np.array([np.nan, 1.0]).mean() + + first, second = compute(), compute() - assert type(compute()) is np.int64 - assert type(compute()) is int + 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): @@ -597,17 +625,77 @@ def test_missing_numpy_is_not_an_unreadable_entry(directory, monkeypatch): same(cache['key'], np.arange(2.0)) -def test_real_memory_shortage_keeps_the_entry(directory, monkeypatch): - def fail(*args, **kwargs): - raise MemoryError +def test_scalar_without_numpy_raises_and_is_kept(directory): + with dc.Cache(directory, disk=dc.SafeArrayDisk) 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.SafeArrayDisk.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) - with dc.Cache(directory, disk=MissDisk) as cache: - cache['key'] = np.arange(2.0) - monkeypatch.setattr(np.lib.format, 'read_array', fail) - with pytest.raises(MemoryError): - cache.get('key') - monkeypatch.undo() - same(cache['key'], np.arange(2.0)) + +def test_arrays_read_back_writeable(directory): + with dc.Cache(directory, disk=dc.SafeArrayDisk) 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=dc.SafeArrayDisk) 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=dc.SafeArrayDisk) as cache: + cache['key'] = array + assert cache['key'].shape == (0,) + + +def test_writing_a_large_array_copies_it_at_most_twice(directory): + array = np.arange(4000000.0) + + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + tracemalloc.start() + cache['key'] = {'array': array} + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert peak < array.nbytes * 2.2 + + +def test_reading_a_large_array_takes_no_second_copy(directory): + array = np.arange(4000000.0) + + with dc.Cache(directory, disk=dc.SafeArrayDisk) 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): @@ -672,16 +760,17 @@ def test_with_types_still_refuses_numpy(): def test_row_format_is_pinned(directory): """Changing what is on disk breaks every cache: this is hand-built.""" array = np.arange(3.0) - write_row(directory, persid(('npy', npy_bytes(array)))) + write_row(directory, persid(array_pid(npy_bytes(array)))) with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: same(cache['key'], array) other = op.join(directory, 'other') - write_row(other, persid(('item', 1.5))) + write_row(other, persid(('scalar', 'float32', 1.5))) with dc.Cache(other, disk=dc.SafeArrayDisk) as cache: assert cache['key'] == 1.5 + assert type(cache['key']) is np.float32 def test_stored_row_uses_the_pinned_format(directory): @@ -694,9 +783,12 @@ def test_stored_row_uses_the_pinned_format(directory): ops = [(op.name, arg) for op, arg, _ in pickletools.genops(rows[0])] assert ('BINPERSID', None) in ops + assert any(name == 'BYTEARRAY8' for name, _ in ops) assert any(arg == 'npy' for _, arg in ops) assert any(arg == npy_bytes(array) for _, arg in ops) - assert any(arg == 'item' for _, arg, _ in pickletools.genops(rows[1])) + args = [arg for _, arg, _ in pickletools.genops(rows[1])] + assert 'scalar' in args + assert 'int64' in args @pytest.mark.parametrize('reader', [dc.Disk, dc.SafeDisk]) @@ -754,7 +846,7 @@ def test_poison_does_not_run(directory): def test_poison_inside_a_persistent_id_does_not_run(directory): marker = directory + '.executed' - write_row(directory, persid(('item', Poison(marker)))) + write_row(directory, persid(('scalar', 'int64', Poison(marker)))) with dc.Cache(directory, disk=MissDisk) as cache: assert cache.get('key', 'default') == 'default' @@ -764,7 +856,7 @@ def test_poison_inside_a_persistent_id_does_not_run(directory): def test_pickle_in_the_array_slot_does_not_run(directory): marker = directory + '.executed' - write_row(directory, persid(('npy', pickle.dumps(Poison(marker))))) + write_row(directory, persid(array_pid(pickle.dumps(Poison(marker))))) with dc.Cache(directory, disk=MissDisk) as cache: assert cache.get('key', 'default') == 'default' @@ -806,7 +898,7 @@ def test_pickle_in_the_array_slot_does_not_run(directory): @pytest.mark.parametrize('data', FORGED.values(), ids=FORGED.keys()) def test_forged_array_is_refused(directory, data): - write_row(directory, persid(('npy', data))) + write_row(directory, persid(array_pid(data))) with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: with pytest.raises(pickle.UnpicklingError): @@ -815,7 +907,7 @@ def test_forged_array_is_refused(directory, data): @pytest.mark.parametrize('data', FORGED.values(), ids=FORGED.keys()) def test_forged_array_is_a_miss(directory, data): - write_row(directory, persid(('npy', data))) + write_row(directory, persid(array_pid(data))) with dc.Cache(directory, disk=MissDisk) as cache: assert cache.get('key', 'default') == 'default' @@ -826,17 +918,28 @@ def test_forged_array_is_a_miss(directory, data): FORGED_IDS = { 'bytes': b'data', - 'none': None, - 'list': ['npy', b'data'], + 'list': ['npy', bytearray(npy_bytes(np.arange(2.0)))], + 'set': frozenset(['npy']), 'one item': ('npy',), - 'three items': ('npy', b'data', 1), - 'unknown kind': ('other', b'data'), - 'kind not a string': (1, b'data'), + 'three items': ('npy', bytearray(b'data'), 1), + 'unknown kind': ('other', bytearray(b'data')), + 'old item kind': ('item', 1.5), + 'kind not a string': (1, bytearray(b'data')), + 'array of bytes': ('npy', npy_bytes(np.arange(2.0))), 'array of str': ('npy', 'text'), 'array of int': ('npy', 1), 'array of none': ('npy', None), 'array of list': ('npy', [1, 2]), - 'array of bytearray': ('npy', bytearray(b'data')), + 'scalar with two items': ('scalar', 'float64'), + 'scalar with four items': ('scalar', 'float64', 1.0, 1), + 'scalar of unknown type': ('scalar', 'object', 1), + 'scalar of a numpy function': ('scalar', 'frombuffer', b'data'), + 'scalar type not a string': ('scalar', ['float64'], 1.0), + 'scalar type number': ('scalar', 1, 1.0), + 'scalar of a list': ('scalar', 'float64', [1.0, 2.0]), + 'scalar of none': ('scalar', 'float64', None), + 'scalar of text': ('scalar', 'float64', 'not a number'), + 'scalar too large': ('scalar', 'int64', 10**30), } @@ -855,8 +958,8 @@ def test_forged_persistent_id_is_refused(directory, pid): @pytest.mark.parametrize( 'pid', - [('npy', npy_bytes(np.arange(2.0))), ('item', 1)], - ids=['array', 'item'], + [array_pid(npy_bytes(np.arange(2.0))), ('scalar', 'int64', 1)], + ids=['array', 'scalar'], ) def test_persistent_id_in_a_key_row_is_refused(directory, pid): with dc.Cache(directory, disk=dc.SafeDisk) as cache: @@ -884,21 +987,37 @@ def test_keys_still_read_back(directory): assert list(cache) == [('a', 1), 'text'] +def test_none_as_a_persistent_id_is_refused(directory): + write_row(directory, b'\x80\x05NQ.') + + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with pytest.raises(pickle.UnpicklingError): + cache['key'] + + +@pytest.mark.parametrize( + 'name', + ['huge shape', 'truncated', 'object', 'object field', 'subarray object'], +) +def test_header_check_refuses_before_numpy_sees_the_data(directory, name): + write_row(directory, persid(array_pid(FORGED[name]))) + + with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with pytest.raises(pickle.UnpicklingError, match='array header'): + cache['key'] + + def test_forged_array_nested_in_a_value_is_a_miss(directory): - inner = persid(('npy', FORGED['huge shape'])) - write_row(directory, b'\x80\x03]' + inner[2:-1] + b'a.') + write_row(directory, forge([Pid(array_pid(FORGED['huge shape']))])) with dc.Cache(directory, disk=MissDisk) as cache: assert cache.get('key', 'default') == 'default' -def test_one_array_referenced_many_times_decodes_once(directory): +def test_one_id_used_many_times_decodes_once(directory): """A row cannot decode to more than the data it holds.""" - array = np.zeros(1000) - pid = pickle.dumps(('npy', npy_bytes(array)), protocol=3)[2:-1] - data = b'\x80\x03]q\x00(' + pid + b'q\x01Q' - data += b'h\x01Q' * 1000 + b'e.' - write_row(directory, data) + pid = Pid(array_pid(npy_bytes(np.zeros(1000)))) + write_row(directory, forge([pid] * 1001)) with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: result = cache['key'] @@ -908,11 +1027,8 @@ def test_one_array_referenced_many_times_decodes_once(directory): def test_one_buffer_in_many_ids_decodes_once(directory): - array = np.zeros(1000) - data = b'\x80\x03]q\x00(\x8c\x03npyq\x01' - data += pickle.dumps(npy_bytes(array), protocol=3)[2:-1] + b'q\x02\x86Q' - data += b'h\x01h\x02\x86Q' * 1000 + b'e.' - write_row(directory, data) + data = bytearray(npy_bytes(np.zeros(1000))) + write_row(directory, forge([Pid(('npy', data)) for _ in range(1001)])) with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: result = cache['key'] @@ -923,8 +1039,8 @@ def test_one_buffer_in_many_ids_decodes_once(directory): def test_deep_nesting_reads_without_recursing(directory): # Built by hand, because pickling this nesting would itself recurse. - array = persid(('npy', npy_bytes(np.zeros(1)))) - data = b'\x80\x03' + b']' * 100000 + array[2:-1] + b'a' * 100000 + b'.' + leaf = pickle.dumps(('scalar', 'float64', 1.0), protocol=3)[2:-1] + data = b'\x80\x03' + b']' * 100000 + leaf + b'Q' + b'a' * 100000 + b'.' write_row(directory, data) with dc.Cache(directory, disk=MissDisk) as cache: 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' From 3514962a0509874c8439889cff56374535bcd241 Mon Sep 17 00:00:00 2001 From: August Vi Date: Mon, 2 Nov 2026 07:07:06 -0800 Subject: [PATCH 07/10] Replace SafeArrayDisk with SafeDisk.with_numpy(), which combines with with_types and with_options in any order. --- CHANGELOG.md | 2 +- CONTRIBUTING.md | 4 +- README.md | 4 +- SECURITY.md | 4 +- diskcache_community/__init__.py | 2 - diskcache_community/__init__.pyi | 2 - diskcache_community/arrays.py | 134 +++++------- diskcache_community/arrays.pyi | 9 +- diskcache_community/core.py | 72 ++++++- diskcache_community/core.pyi | 8 +- docs/api.rst | 10 - docs/fix-cve-2025-69872.rst | 16 +- docs/tutorial.rst | 52 ++--- ...afearraydisk.py => test_safedisk_numpy.py} | 191 +++++++++++------- 14 files changed, 287 insertions(+), 223 deletions(-) rename tests/{test_safearraydisk.py => test_safedisk_numpy.py} (83%) diff --git a/CHANGELOG.md b/CHANGELOG.md index b9462ea..61ba1f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ 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). -- `SafeArrayDisk` is a `SafeDisk` 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`. 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 or by 5.6.3, which raise `UnpicklingError`. numpy stays optional, and `SafeDisk` still refuses it (PR #34). +- `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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3a3debe..ddf1179 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -46,7 +46,7 @@ This fork publishes the same library under a new distribution name, `diskcache-c - 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 row format `SafeArrayDisk` writes: a persistent ID of `('npy', )` or `('scalar', , )` in a pickle of protocol 5 or higher, with `tests/test_safearraydisk.py` pinning it; +- the row format `SafeDisk.with_numpy()` writes: a persistent ID of `('npy', )` or `('scalar', , )` in a pickle of protocol 5 or higher, with `tests/test_safedisk_numpy.py` pinning it; - 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. Entries a process wrote through `SafeArrayDisk` that hold numpy data are the one exception to reading them back: 5.6.3 opens the directory and raises `UnpicklingError` for those entries. Golden fixtures written by the 5.6.3 wheel guard this in the test suite. +A process running 5.6.3 must be able to open a directory written by this fork, and the reverse. Entries a process wrote through `SafeDisk.with_numpy()` that hold numpy data are the one exception to reading them back: 5.6.3 opens the directory and raises `UnpicklingError` for those entries. Golden fixtures written by the 5.6.3 wheel guard this in the test suite. diff --git a/README.md b/README.md index a20f507..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`, `SafeJSONDisk` and `SafeArrayDisk` 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 42ab948..b0549da 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -19,10 +19,10 @@ 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. -- `SafeArrayDisk` is 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. +- `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`, `SafeJSONDisk` or `SafeArrayDisk`, 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. +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. diff --git a/diskcache_community/__init__.py b/diskcache_community/__init__.py index 4c7f1b1..d97178a 100644 --- a/diskcache_community/__init__.py +++ b/diskcache_community/__init__.py @@ -5,7 +5,6 @@ The :doc:`tutorial` provides a helpful walkthrough of most methods. """ -from .arrays import SafeArrayDisk from .core import ( DEFAULT_SETTINGS, ENOVAL, @@ -51,7 +50,6 @@ 'Lock', 'RLock', 'ReadOnlyError', - 'SafeArrayDisk', 'SafeDisk', 'SafeJSONDisk', 'Timeout', diff --git a/diskcache_community/__init__.pyi b/diskcache_community/__init__.pyi index 977da3f..b3f34e5 100644 --- a/diskcache_community/__init__.pyi +++ b/diskcache_community/__init__.pyi @@ -1,4 +1,3 @@ -from diskcache_community.arrays import SafeArrayDisk as SafeArrayDisk from diskcache_community.core import DEFAULT_SETTINGS as DEFAULT_SETTINGS from diskcache_community.core import ENOVAL as ENOVAL from diskcache_community.core import EVICTION_POLICY as EVICTION_POLICY @@ -41,7 +40,6 @@ __all__ = [ 'Lock', 'RLock', 'ReadOnlyError', - 'SafeArrayDisk', 'SafeDisk', 'SafeJSONDisk', 'Timeout', diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py index 4ed8709..bb212e4 100644 --- a/diskcache_community/arrays.py +++ b/diskcache_community/arrays.py @@ -3,16 +3,8 @@ import io import math import pickle -import sys -from .core import ( - SAFE_PICKLE_NAMES, - UNKNOWN, - Disk, - SafeDisk, - SafeUnpickler, - _ReaderMissing, -) +from .core import SAFE_PICKLE_NAMES, Disk, SafeUnpickler, _ReaderMissing # Pickle protocol 5 is the first that stores a bytearray in place, so that # reading a file-backed entry needs no second copy of its array. @@ -51,7 +43,7 @@ def _import_numpy(): import numpy except ImportError as error: raise _ReaderMissing( - 'SafeArrayDisk needs numpy to read this entry' + 'SafeDisk needs numpy to read this entry' ) from error return numpy @@ -78,7 +70,7 @@ def _read_header(numpy, data): raise ValueError('unsupported .npy version %r' % (version,)) except Exception as error: raise pickle.UnpicklingError( - 'SafeArrayDisk refuses the array header' + 'SafeDisk refuses the array header' ) from error shape, fortran, dtype = header @@ -91,7 +83,7 @@ def _read_header(numpy, data): or dtype.itemsize * count != len(data) - offset or (count and not dtype.itemsize) ): - raise pickle.UnpicklingError('SafeArrayDisk refuses the array header') + raise pickle.UnpicklingError('SafeDisk refuses the array header') return shape, fortran, dtype, offset @@ -109,7 +101,7 @@ def _pack_array(numpy, array): _read_header(numpy, view) except (ValueError, pickle.UnpicklingError) as error: raise TypeError( - 'SafeArrayDisk cannot store an array of dtype %s' % (array.dtype,) + 'SafeDisk cannot store an array of dtype %s' % (array.dtype,) ) from error return pickle.PickleBuffer(view) @@ -132,7 +124,7 @@ def _unpack_array(data): return array.reshape(shape, order='F' if fortran else 'C') except Exception as error: raise pickle.UnpicklingError( - 'SafeArrayDisk refuses the array data' + 'SafeDisk refuses the array data' ) from error @@ -142,7 +134,7 @@ def _scalar_name(scalar): name = {'U': 'str', 'S': 'bytes'}.get(dtype.kind, dtype.name) if name not in _SCALARS: raise TypeError( - 'SafeArrayDisk cannot store a scalar of type numpy.%s' + 'SafeDisk cannot store a scalar of type numpy.%s' % (type(scalar).__name__,) ) return name @@ -152,16 +144,14 @@ def _unpack_scalar(name, value): """Return the numpy scalar `value` of type `name`, which may be forged.""" # pylint: disable=broad-except,unidiomatic-typecheck if type(value) not in _SCALAR_VALUES: - raise pickle.UnpicklingError('SafeArrayDisk refuses a scalar') + raise pickle.UnpicklingError('SafeDisk refuses a scalar') numpy = _import_numpy() try: return getattr(numpy, _SCALARS[name])(value) except Exception as error: - raise pickle.UnpicklingError( - 'SafeArrayDisk refuses a scalar' - ) from error + raise pickle.UnpicklingError('SafeDisk refuses a scalar') from error class _ArrayPickler(pickle.Pickler): @@ -186,7 +176,7 @@ def persistent_id(self, obj): if isinstance(obj, numpy.ndarray): if type(obj) is not numpy.ndarray: raise TypeError( - 'SafeArrayDisk cannot store a subclass of ndarray: %s.%s' + 'SafeDisk cannot store a subclass of ndarray: %s.%s' % (type(obj).__module__, type(obj).__qualname__) ) # Holding `obj` keeps its id from being reused: pickle's own memo @@ -204,7 +194,7 @@ def persistent_id(self, obj): value = obj.item() if isinstance(value, numpy.generic): raise TypeError( - 'SafeArrayDisk cannot store a scalar of type numpy.%s' + 'SafeDisk cannot store a scalar of type numpy.%s' % (kind.__name__,) ) return ('scalar', name, value) @@ -234,82 +224,50 @@ def persistent_load(self, pid): entry = (payload, _unpack_array(payload)) self._arrays[id(payload)] = entry return entry[1] - raise pickle.UnpicklingError('SafeArrayDisk refuses a persistent ID') + raise pickle.UnpicklingError('SafeDisk refuses a persistent ID') -class SafeArrayDisk(SafeDisk): - """Disk that stores numpy arrays as ``.npy`` bytes instead of pickles. +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, + ) - :class:`SafeDisk` refuses numpy, because unpickling an array is not - memory-safe (see ``UNSAFE_MODULES``). :class:`SafeArrayDisk` keeps those - pickles out of the cache by writing arrays through pickle's persistent IDs - in numpy's own format: - * An array anywhere in a value is saved with ``allow_pickle=False``, and - its header is checked before it is read back. That includes the - containers :class:`SafeDisk` accepts and the classes allowed with - :meth:`SafeDisk.with_types`. - * numpy booleans, integers, floats, complex numbers and strings read back - as the same numpy type. - * Arrays of objects, subclasses of ``ndarray`` such as masked arrays, and - other numpy scalars raise :exc:`TypeError` when written. +def store(disk, numpy, value, key): + """Return the Cache table fields of `value`, which may hold numpy data. - Keys and everything else behave as in :class:`SafeDisk`. numpy stays - optional: it is used only for a value that needs it, and a process without - it raises :exc:`ImportError` for an entry holding numpy data and keeps the - entry, even with ``on_error='miss'``. - - 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. Every process that shares a directory should select this disk. + `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) - def _unpickler(self, file): - return _ArrayUnpickler( - file, - SAFE_PICKLE_NAMES, - self.safe_types, - self.safe_modules, - ) - - def get(self, key, raw): - if raw: - return super().get(key, raw) - # Keys never hold numpy, so a key row gets no persistent IDs. - return super()._unpickler(io.BytesIO(key)).load() - - def store(self, value, read, key=UNKNOWN): - # pylint: disable=unidiomatic-typecheck - numpy = sys.modules.get('numpy') - - if numpy is None or read or type(value) in (bytes, str, int, float): - return super().store(value, read, key=key) - - try: - self._check(value, 'value') - except TypeError: - pass - else: - return Disk.store(self, value, read, key=key) - - buffer = io.BytesIO() - protocol = max(self.pickle_protocol, _MIN_PROTOCOL) - pickler = _ArrayPickler(buffer, protocol, numpy) + 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( - 'SafeArrayDisk cannot store value: %s' % (error,) - ) from error + try: + pickler.dump(value) + except (pickle.PicklingError, AttributeError) as error: + raise TypeError( + 'SafeDisk cannot store value: %s' % (error,) + ) from error - data = buffer.getvalue() + 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: - self._check_pickle(data, value, 'value') + # 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 self._store_pickle(data, value, key) + return disk._store_pickle(data, value, key) diff --git a/diskcache_community/arrays.pyi b/diskcache_community/arrays.pyi index ae596ce..84a46f0 100644 --- a/diskcache_community/arrays.pyi +++ b/diskcache_community/arrays.pyi @@ -1,3 +1,8 @@ -from diskcache_community.core import SafeDisk +from typing import IO, Any -class SafeArrayDisk(SafeDisk): ... +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 df5cdd4..1e1d887 100644 --- a/diskcache_community/core.py +++ b/diskcache_community/core.py @@ -805,6 +805,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) @@ -826,8 +827,7 @@ def with_types(cls, *types, modules=()): 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``; use - :class:`SafeArrayDisk ` to store - arrays. + :meth:`with_numpy` to store arrays. >>> import collections >>> Point = collections.namedtuple('Point', 'x y') @@ -848,7 +848,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,) @@ -864,7 +867,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( @@ -877,7 +883,36 @@ def with_types(cls, *types, modules=()): }, ) - def _unpickler(self, file): + @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( file, SAFE_PICKLE_NAMES, @@ -885,6 +920,11 @@ def _unpickler(self, file): 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) @@ -1014,12 +1054,17 @@ def put(self, key): def get(self, key, raw): if raw: return super().get(key, raw) - return self._unpickler(io.BytesIO(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): @@ -1055,10 +1100,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') @@ -3943,3 +3992,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/api.rst b/docs/api.rst index 236682e..85c419b 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -167,16 +167,6 @@ Read the :ref:`SafeJSONDisk tutorial ` for details. :special-members: :exclude-members: __weakref__ -SafeArrayDisk -------------- - -Read the :ref:`SafeArrayDisk tutorial ` for details. - -.. autoclass:: diskcache_community.SafeArrayDisk - :members: - :special-members: - :exclude-members: __weakref__ - Default disk ------------ diff --git a/docs/fix-cve-2025-69872.rst b/docs/fix-cve-2025-69872.rst index 2dd18d4..78bca96 100644 --- a/docs/fix-cve-2025-69872.rst +++ b/docs/fix-cve-2025-69872.rst @@ -111,19 +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. Select - :class:`SafeArrayDisk ` to store arrays - instead, which saves them in numpy's ``.npy`` format: + 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 - from diskcache_community import Cache, SafeArrayDisk + from diskcache_community import Cache, SafeDisk - cache = Cache('/var/tmp/mycache', disk=SafeArrayDisk) + ArrayDisk = SafeDisk.with_numpy().with_types(Message, Point) + cache = Cache('/var/tmp/mycache', disk=ArrayDisk) cache['embedding'] = array - In the environment, set - ``DISKCACHE_DISK=diskcache_community:SafeArrayDisk``. The same lists apply. + 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 b92a96a..1a372e8 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. Use :class:`SafeArrayDisk -` 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. +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,31 +1056,34 @@ extend and :meth:`SafeDisk.with_types ``SafeJSONDisk.with_options(on_error='miss')`` to drop refused entries, as described in :ref:`tutorial-unreadable-entries`. -.. _tutorial-safearraydisk: +.. _tutorial-with-numpy: -SafeArrayDisk -............. +Storing numpy arrays +.................... :class:`SafeDisk ` refuses numpy, because -unpickling an array is not memory-safe. :class:`SafeArrayDisk -` stores arrays without pickling them, by -writing them through pickle's persistent IDs in numpy's ``.npy`` format: +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, SafeArrayDisk + from diskcache_community import Cache, SafeDisk - cache = Cache('/tmp/mycache', disk=SafeArrayDisk) + 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 saved with ``allow_pickle=False``, including inside classes allowed with :meth:`SafeDisk.with_types - `. Its header is checked against the - data before it is read back, also with ``allow_pickle=False``. An array that - appears twice in a value is stored once and is one array when read back. + `, which can be called before or + after ``with_numpy``. Its header is checked against the data before it is + read back, also with ``allow_pickle=False``. An array that appears twice in a + value is stored once and is one array when read back. * 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 @@ -1099,14 +1102,17 @@ the same way. An entry whose array data has been tampered with raises ``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, or by diskcache 5.6.3: they open the +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. -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. +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: @@ -1134,8 +1140,8 @@ or in the environment, without changing any code: ``DISKCACHE_DISK`` takes ``safe``, ``default``, or the import path of a disk class such as ``myapp.cache:AppDisk``. With ``safe`` or a subclass of -:class:`SafeDisk `, such as -``diskcache_community:SafeArrayDisk``, ``DISKCACHE_SAFE_TYPES`` and +: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 diff --git a/tests/test_safearraydisk.py b/tests/test_safedisk_numpy.py similarity index 83% rename from tests/test_safearraydisk.py rename to tests/test_safedisk_numpy.py index dfc00a6..5ca75d7 100644 --- a/tests/test_safearraydisk.py +++ b/tests/test_safedisk_numpy.py @@ -1,4 +1,4 @@ -"""Test diskcache_community.SafeArrayDisk.""" +"""Test SafeDisk.with_numpy().""" import collections as co import dataclasses @@ -17,14 +17,15 @@ import pytest import diskcache_community as dc -from diskcache_community import arrays +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) -MissDisk = dc.SafeArrayDisk.with_options(on_error='miss') +ArrayDisk = dc.SafeDisk.with_numpy() +MissDisk = ArrayDisk.with_options(on_error='miss') def _execute(path): @@ -180,7 +181,7 @@ def npy_header(header, body=b'', version=1): ids=lambda array: '%s-%s' % (array.dtype, array.shape), ) def test_array_round_trip(directory, array): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = array same(cache['key'], array) @@ -189,7 +190,7 @@ def test_arrays_in_containers(directory): array = np.arange(3.0) value = {'a': array, 'b': [array, (array, 1)], 'c': 'text'} - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = value result = cache['key'] @@ -215,7 +216,7 @@ def test_arrays_in_the_other_safe_containers(directory, make): array = np.arange(2.0) value = make(array) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = value result = cache['key'] @@ -226,7 +227,7 @@ def test_arrays_in_the_other_safe_containers(directory, make): def test_arrays_in_allowed_classes(directory): array = np.arange(3.0) - disk = dc.SafeArrayDisk.with_types(Result, Pair, State) + disk = ArrayDisk.with_types(Result, Pair, State) with dc.Cache(directory, disk=disk) as cache: cache['dataclass'] = Result('name', array) @@ -239,7 +240,7 @@ def test_arrays_in_allowed_classes(directory): def test_arrays_in_classes_outside_the_allowlist_are_refused(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(TypeError): cache['key'] = Result('name', np.arange(2.0)) with pytest.raises(TypeError): @@ -250,7 +251,7 @@ def test_writing_leaves_the_caller_value_alone(directory): array = np.arange(3.0) value = {'a': array} - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = value assert value['a'] is array @@ -259,7 +260,7 @@ def test_writing_leaves_the_caller_value_alone(directory): def test_shared_arrays_stay_shared_and_are_stored_once(directory): array = np.arange(1000.0) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['one'] = [array] cache['many'] = [array] * 100 result = cache['many'] @@ -270,7 +271,7 @@ def test_shared_arrays_stay_shared_and_are_stored_once(directory): def test_temporary_arrays_keep_their_own_data(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = Temporaries(50) result = cache['key'] @@ -284,7 +285,7 @@ def test_cycles_work(directory): value = [array] value.append(value) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = value result = cache['key'] @@ -298,7 +299,7 @@ def test_shared_structure_stays_linear(directory): for _ in range(64): value = [value, value] - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = value result = cache['key'] @@ -311,7 +312,7 @@ def test_every_pickle_protocol_setting(directory, protocol): pytest.skip('protocol not available') with dc.Cache( - directory, disk=dc.SafeArrayDisk, disk_pickle_protocol=protocol + 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)) @@ -341,7 +342,7 @@ def test_every_pickle_protocol_setting(directory, protocol): ids=lambda value: type(value).__name__, ) def test_scalars_keep_their_type_and_value(directory, scalar): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = scalar cache['nested'] = {'metrics': [scalar, scalar]} assert cache['key'] == scalar @@ -351,7 +352,7 @@ def test_scalars_keep_their_type_and_value(directory, scalar): @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=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['nan'] = kind('nan') cache['inf'] = kind('inf') cache['nested'] = [kind('nan')] @@ -362,7 +363,7 @@ def test_scalar_nan_and_infinity_are_kept(directory, kind): def test_scalars_as_dict_keys_keep_their_type(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + 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 @@ -383,7 +384,7 @@ def test_scalars_as_dict_keys_keep_their_type(directory): def test_scalars_in_the_other_safe_containers(directory, make): value = make(np.int64(2)) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = value result = cache['key'] @@ -402,7 +403,7 @@ def test_scalars_in_the_other_safe_containers(directory, make): ids=lambda value: type(value).__name__, ) def test_time_scalars_are_refused(directory, scalar): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(TypeError): cache['key'] = scalar with pytest.raises(TypeError): @@ -412,7 +413,7 @@ def test_time_scalars_are_refused(directory, scalar): def test_long_double_reads_back_or_is_refused(directory): value = np.longdouble(1.5) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: try: cache['key'] = {'a': value} except TypeError: @@ -430,7 +431,7 @@ def test_long_double_reads_back_or_is_refused(directory): ids=['object', 'field', 'subarray'], ) def test_object_arrays_are_refused(directory, array): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(TypeError): cache['key'] = array assert len(cache) == 0 @@ -471,7 +472,7 @@ def test_arrays_the_reader_would_refuse_fail_at_write(directory, array): ) def test_subclasses_are_refused(directory, make): """numpy.save would silently drop what makes them a subclass.""" - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(TypeError, match='subclass'): cache['key'] = make() with pytest.raises(TypeError, match='subclass'): @@ -479,7 +480,7 @@ def test_subclasses_are_refused(directory, make): def test_numpy_in_a_key_is_refused(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(TypeError): cache[(np.int64(1),)] = 'value' with pytest.raises(TypeError): @@ -489,7 +490,7 @@ def test_numpy_in_a_key_is_refused(directory): def test_memoize(directory): calls = [] - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: @cache.memoize() def compute(size): @@ -503,7 +504,7 @@ def compute(size): def test_memoize_scalar_keeps_its_type_on_a_hit(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: @cache.memoize() def compute(): @@ -516,7 +517,7 @@ def compute(): def test_memoize_scalar_argument_is_refused(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: @cache.memoize() def compute(value): @@ -527,7 +528,7 @@ def compute(value): def test_plain_values_behave_as_in_safedisk(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + 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 @@ -538,7 +539,7 @@ def test_plain_values_behave_as_in_safedisk(directory): def test_plain_values_are_ordinary_pickles(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + 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')) @@ -558,7 +559,7 @@ def refuse(*args): monkeypatch.setattr(arrays, '_ArrayPickler', refuse) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = {'a': [1, 2.5, 'x']} assert cache['key'] == {'a': [1, 2.5, 'x']} @@ -568,7 +569,8 @@ def test_plain_values_do_not_import_numpy(directory): 'import sys\n' 'import diskcache_community as dc\n' 'assert "numpy" not in sys.modules\n' - 'with dc.Cache(sys.argv[1], disk=dc.SafeArrayDisk) as cache:\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' @@ -577,28 +579,29 @@ def test_plain_values_do_not_import_numpy(directory): def test_reads_an_array_in_a_process_that_has_not_imported_numpy(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + 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' - 'with dc.Cache(sys.argv[1], disk=dc.SafeArrayDisk) as cache:\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=dc.SafeArrayDisk) as cache: + 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.SafeArrayDisk.with_options(on_error="miss")\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' @@ -611,7 +614,7 @@ def test_without_numpy_an_entry_raises_and_is_kept(directory): ) subprocess.run([sys.executable, '-c', code, directory], check=True) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: same(cache['key'], np.arange(2.0)) @@ -626,14 +629,14 @@ def test_missing_numpy_is_not_an_unreadable_entry(directory, monkeypatch): def test_scalar_without_numpy_raises_and_is_kept(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + 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.SafeArrayDisk.with_options(on_error="miss")\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' @@ -647,7 +650,7 @@ def test_scalar_without_numpy_raises_and_is_kept(directory): def test_arrays_read_back_writeable(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = np.arange(3.0) result = cache['key'] @@ -659,7 +662,7 @@ def test_arrays_read_back_writeable(directory): def test_fortran_order_round_trip(directory): array = np.arange(6.0).reshape(2, 3, order='F') - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = array assert cache['key'].tolist() == array.tolist() @@ -667,7 +670,7 @@ def test_fortran_order_round_trip(directory): def test_empty_array_of_a_dtype_with_no_width(directory): array = np.zeros(0, dtype=[]) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = array assert cache['key'].shape == (0,) @@ -675,7 +678,7 @@ def test_empty_array_of_a_dtype_with_no_width(directory): def test_writing_a_large_array_copies_it_at_most_twice(directory): array = np.arange(4000000.0) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: tracemalloc.start() cache['key'] = {'array': array} _, peak = tracemalloc.get_traced_memory() @@ -687,7 +690,7 @@ def test_writing_a_large_array_copies_it_at_most_twice(directory): def test_reading_a_large_array_takes_no_second_copy(directory): array = np.arange(4000000.0) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['key'] = array tracemalloc.start() result = cache['key'] @@ -699,7 +702,7 @@ def test_reading_a_large_array_takes_no_second_copy(directory): def test_read_true_still_returns_a_file(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + 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' @@ -708,9 +711,7 @@ def test_read_true_still_returns_a_file(directory): def test_large_array_lives_in_a_file(directory): array = np.arange(100000.0) - with dc.Cache( - directory, disk=dc.SafeArrayDisk, disk_min_file_size=0 - ) as cache: + with dc.Cache(directory, disk=ArrayDisk, disk_min_file_size=0) as cache: cache['key'] = array same(cache['key'], array) assert cache.check() == [] @@ -719,27 +720,23 @@ def test_large_array_lives_in_a_file(directory): def test_fanout_index_and_deque(directory): array = np.arange(3.0) - with dc.FanoutCache( - op.join(directory, 'fanout'), disk=dc.SafeArrayDisk - ) as cache: + 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=dc.SafeArrayDisk) + 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=dc.SafeArrayDisk - ) + 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 = dc.SafeArrayDisk.with_types(Pair).with_options(on_error='miss') + disk = ArrayDisk.with_types(Pair).with_options(on_error='miss') with dc.Cache(directory, disk=disk) as cache: cache['pair'] = Pair(1, 2) @@ -747,14 +744,64 @@ def test_with_types_and_options_combine(directory): assert cache['pair'] == Pair(1, 2) same(cache['array'], np.arange(2.0)) copy = pickle.loads(pickle.dumps(cache)) - assert isinstance(copy.disk, dc.SafeArrayDisk) + 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): - dc.SafeArrayDisk.with_types(np.ndarray) + 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): @@ -762,13 +809,13 @@ def test_row_format_is_pinned(directory): array = np.arange(3.0) write_row(directory, persid(array_pid(npy_bytes(array)))) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: same(cache['key'], array) other = op.join(directory, 'other') write_row(other, persid(('scalar', 'float32', 1.5))) - with dc.Cache(other, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(other, disk=ArrayDisk) as cache: assert cache['key'] == 1.5 assert type(cache['key']) is np.float32 @@ -776,7 +823,7 @@ def test_row_format_is_pinned(directory): def test_stored_row_uses_the_pinned_format(directory): array = np.arange(3.0) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['array'] = array cache['scalar'] = [np.int64(1)] rows = row_bytes(cache, 'array'), row_bytes(cache, 'scalar') @@ -793,7 +840,7 @@ def test_stored_row_uses_the_pinned_format(directory): @pytest.mark.parametrize('reader', [dc.Disk, dc.SafeDisk]) def test_another_disk_refuses_an_array_entry(directory, reader): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['array'] = np.arange(2.0) cache['plain'] = {'a': 1} @@ -804,7 +851,7 @@ def test_another_disk_refuses_an_array_entry(directory, reader): def test_another_disk_with_miss_deletes_an_array_entry(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache['array'] = np.arange(2.0) disk = dc.SafeDisk.with_options(on_error='miss') @@ -818,7 +865,7 @@ def test_reads_safedisk_entries(directory): with dc.Cache(directory, disk=dc.SafeDisk) as cache: cache['key'] = {'a': [1, 2]} - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: assert cache['key'] == {'a': [1, 2]} @@ -826,7 +873,7 @@ def test_array_pickled_by_disk_is_refused(directory): with dc.Cache(directory, disk=dc.Disk) as cache: cache['key'] = np.arange(2.0) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(pickle.UnpicklingError): cache['key'] @@ -900,7 +947,7 @@ def test_pickle_in_the_array_slot_does_not_run(directory): def test_forged_array_is_refused(directory, data): write_row(directory, persid(array_pid(data))) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(pickle.UnpicklingError): cache['key'] @@ -947,7 +994,7 @@ def test_forged_array_is_a_miss(directory, data): def test_forged_persistent_id_is_refused(directory, pid): write_row(directory, persid(pid)) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(pickle.UnpicklingError): cache['key'] @@ -972,7 +1019,7 @@ def test_persistent_id_in_a_key_row_is_refused(directory, pid): con.commit() con.close() - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(pickle.UnpicklingError): list(cache) @@ -981,7 +1028,7 @@ def test_persistent_id_in_a_key_row_is_refused(directory, pid): def test_keys_still_read_back(directory): - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: cache[('a', 1)] = np.arange(2.0) cache['text'] = np.arange(2.0) assert list(cache) == [('a', 1), 'text'] @@ -990,7 +1037,7 @@ def test_keys_still_read_back(directory): def test_none_as_a_persistent_id_is_refused(directory): write_row(directory, b'\x80\x05NQ.') - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(pickle.UnpicklingError): cache['key'] @@ -1002,7 +1049,7 @@ def test_none_as_a_persistent_id_is_refused(directory): def test_header_check_refuses_before_numpy_sees_the_data(directory, name): write_row(directory, persid(array_pid(FORGED[name]))) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: with pytest.raises(pickle.UnpicklingError, match='array header'): cache['key'] @@ -1019,7 +1066,7 @@ def test_one_id_used_many_times_decodes_once(directory): pid = Pid(array_pid(npy_bytes(np.zeros(1000)))) write_row(directory, forge([pid] * 1001)) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: result = cache['key'] assert len(result) == 1001 @@ -1030,7 +1077,7 @@ def test_one_buffer_in_many_ids_decodes_once(directory): data = bytearray(npy_bytes(np.zeros(1000))) write_row(directory, forge([Pid(('npy', data)) for _ in range(1001)])) - with dc.Cache(directory, disk=dc.SafeArrayDisk) as cache: + with dc.Cache(directory, disk=ArrayDisk) as cache: result = cache['key'] assert len(result) == 1001 From d51ccd66fd3fd92af0d394bd548b2467c324523a Mon Sep 17 00:00:00 2001 From: August Vi Date: Mon, 2 Nov 2026 08:30:30 -0800 Subject: [PATCH 08/10] Decode each scalar string once, require scalar values to match their type name, and read past large strings and buffers in the write-time pickle scan. --- .pylintrc | 2 +- diskcache_community/arrays.py | 61 ++++++++++++++-------- diskcache_community/core.py | 45 +++++++++++++++- tests/test_safedisk_numpy.py | 96 +++++++++++++++++++++++++++++++++++ tests/test_safedisk_types.py | 35 +++++++++++++ 5 files changed, 215 insertions(+), 24 deletions(-) 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/diskcache_community/arrays.py b/diskcache_community/arrays.py index bb212e4..96cb699 100644 --- a/diskcache_community/arrays.py +++ b/diskcache_community/arrays.py @@ -14,27 +14,27 @@ # longest header that numpy parses by default, plus padding. _HEADER_LIMIT = 12 + 10000 + 64 -# The numpy scalar types that are stored, by the name written to the row. The +# 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_', - 'int8': 'int8', - 'int16': 'int16', - 'int32': 'int32', - 'int64': 'int64', - 'uint8': 'uint8', - 'uint16': 'uint16', - 'uint32': 'uint32', - 'uint64': 'uint64', - 'float16': 'float16', - 'float32': 'float32', - 'float64': 'float64', - 'complex64': 'complex64', - 'complex128': 'complex128', - 'str': 'str_', - 'bytes': 'bytes_', + '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), } -_SCALAR_VALUES = (bool, int, float, complex, str, bytes) def _import_numpy(): @@ -141,15 +141,22 @@ def _scalar_name(scalar): def _unpack_scalar(name, value): - """Return the numpy scalar `value` of type `name`, which may be forged.""" + """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 - if type(value) not in _SCALAR_VALUES: + 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, _SCALARS[name])(value) + return getattr(numpy, entry[0])(value) except Exception as error: raise pickle.UnpicklingError('SafeDisk refuses a scalar') from error @@ -207,12 +214,22 @@ class _ArrayUnpickler(SafeUnpickler): 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: if len(pid) == 3 and pid[0] == 'scalar': - return _unpack_scalar(pid[1], pid[2]) + 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 len(pid) == 2 and pid[0] == 'npy': payload = pid[1] if type(payload) is bytearray: diff --git a/diskcache_community/core.py b/diskcache_community/core.py index 1e1d887..e1b7944 100644 --- a/diskcache_community/core.py +++ b/diskcache_community/core.py @@ -610,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.""" @@ -631,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: diff --git a/tests/test_safedisk_numpy.py b/tests/test_safedisk_numpy.py index 5ca75d7..022fb83 100644 --- a/tests/test_safedisk_numpy.py +++ b/tests/test_safedisk_numpy.py @@ -350,6 +350,28 @@ def test_scalars_keep_their_type_and_value(directory, 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: @@ -687,6 +709,22 @@ def test_writing_a_large_array_copies_it_at_most_twice(directory): assert peak < array.nbytes * 2.2 +def test_writing_an_array_in_an_allowed_class_copies_it_at_most_twice( + 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 * 2.2 + + def test_reading_a_large_array_takes_no_second_copy(directory): array = np.arange(4000000.0) @@ -987,6 +1025,16 @@ def test_forged_array_is_a_miss(directory, data): 'scalar of none': ('scalar', 'float64', None), 'scalar of text': ('scalar', 'float64', 'not a number'), 'scalar too large': ('scalar', 'int64', 10**30), + 'bytes of an int': ('scalar', 'bytes', 5), + 'str of an int': ('scalar', 'str', 5), + 'bool of a str': ('scalar', 'bool', 'False'), + 'bool of an int': ('scalar', 'bool', 1), + 'int of a bool': ('scalar', 'int64', True), + 'int of a float': ('scalar', 'int64', 1.0), + 'float of an int': ('scalar', 'float64', 1), + 'complex of a float': ('scalar', 'complex128', 1.0), + 'str of bytes': ('scalar', 'str', b'text'), + 'bytes of a str': ('scalar', 'bytes', 'text'), } @@ -1084,6 +1132,54 @@ def test_one_buffer_in_many_ids_decodes_once(directory): assert all(item is result[0] for item in result) +@pytest.mark.parametrize( + 'name, value', + [('str', 'x' * 100000), ('bytes', b'x' * 100000)], + ids=['str', 'bytes'], +) +def test_one_string_in_many_scalar_ids_decodes_once(directory, name, value): + """Scalar strings sharing one memoized value are copied once.""" + write_row( + directory, forge([Pid(('scalar', name, value)) for _ in range(1001)]) + ) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + tracemalloc.start() + result = cache['key'] + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert len(result) == 1001 + assert all(item is result[0] for item in result) + assert type(result[0]) is getattr(np, name + '_') + assert result[0] == value + assert peak < len(value) * 20 + + +def test_unshared_scalar_strings_stay_distinct(directory): + """A freed string's id may be reused, so the memo holds the string.""" + # Built by hand, so that no string is memoized by the pickle itself. + items = b'' + for index in range(100): + text = ('%04d' % index).encode() * 300 + items += ( + b'\x8c\x06scalar' + + b'\x8c\x03str' + + b'X' + + struct.pack(' Date: Mon, 2 Nov 2026 09:46:58 -0800 Subject: [PATCH 09/10] Store an array's own buffer in the row next to its .npy header, so writing copies it once, and refuse BUILD on a stored array by test. --- CONTRIBUTING.md | 2 +- diskcache_community/arrays.py | 133 +++++++++------ docs/tutorial.rst | 15 +- tests/test_safedisk_numpy.py | 293 +++++++++++++++++++++++++++++----- 4 files changed, 347 insertions(+), 96 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ddf1179..c1365b4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -46,7 +46,7 @@ This fork publishes the same library under a new distribution name, `diskcache-c - 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 row format `SafeDisk.with_numpy()` writes: a persistent ID of `('npy', )` or `('scalar', , )` in a pickle of protocol 5 or higher, with `tests/test_safedisk_numpy.py` pinning it; +- the row format `SafeDisk.with_numpy()` writes: a persistent ID of `('npy', , )` or `('scalar', , )` in a pickle of protocol 5 or higher, with `tests/test_safedisk_numpy.py` pinning it; - 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. Entries a process wrote through `SafeDisk.with_numpy()` that hold numpy data are the one exception to reading them back: 5.6.3 opens the directory and raises `UnpicklingError` for those entries. Golden fixtures written by the 5.6.3 wheel guard this in the test suite. diff --git a/diskcache_community/arrays.py b/diskcache_community/arrays.py index 96cb699..f3beab4 100644 --- a/diskcache_community/arrays.py +++ b/diskcache_community/arrays.py @@ -3,17 +3,15 @@ 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 bytearray in place, so that -# reading a file-backed entry needs no second copy of its array. +# 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 -# Bytes of a ``.npy`` file that can hold its header: the prefix, plus the -# longest header that numpy parses by default, plus padding. -_HEADER_LIMIT = 12 + 10000 + 64 - # 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. @@ -48,24 +46,25 @@ def _import_numpy(): return numpy -def _read_header(numpy, data): - """Return shape, Fortran order, dtype and offset of ``.npy`` bytes `data`. +def _read_header(numpy, header, size): + """Return shape, Fortran order and dtype from the ``.npy`` `header`. - Raise UnpicklingError unless the header is consistent with the 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. + 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(bytes(memoryview(data)[:_HEADER_LIMIT])) + reader = io.BytesIO(header) version = numpy.lib.format.read_magic(reader) if version == (1, 0): - header = numpy.lib.format.read_array_header_1_0(reader) + fields = numpy.lib.format.read_array_header_1_0(reader) elif version == (2, 0): - header = numpy.lib.format.read_array_header_2_0(reader) + fields = numpy.lib.format.read_array_header_2_0(reader) else: raise ValueError('unsupported .npy version %r' % (version,)) except Exception as error: @@ -73,54 +72,70 @@ def _read_header(numpy, data): 'SafeDisk refuses the array header' ) from error - shape, fortran, dtype = header - offset = reader.tell() + 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 dtype.itemsize * count != len(data) - offset + 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, offset + return shape, fortran, dtype def _pack_array(numpy, array): - """Return a buffer of the ``.npy`` bytes of `array`, for pickle to store. + """Return the ``.npy`` header of `array` and a buffer of its data. - The buffer is writable, so pickle stores it as a bytearray in place. + 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: - numpy.save(buffer, array, allow_pickle=False) - view = buffer.getbuffer() - _read_header(numpy, view) + 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 pickle.PickleBuffer(view) + + return header, pickle.PickleBuffer(array) -def _unpack_array(data): +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. `data` may have been tampered with. + `data` is. `header` and `data` may have been tampered with. """ # pylint: disable=broad-except numpy = _import_numpy() - shape, fortran, dtype, offset = _read_header(numpy, data) + 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, offset=offset) + array = numpy.frombuffer(data, dtype) return array.reshape(shape, order='F' if fortran else 'C') except Exception as error: raise pickle.UnpicklingError( @@ -172,7 +187,14 @@ def __init__(self, file, protocol, numpy): self.reduced = False def reducer_override(self, obj): - """Note that an object was reduced, which can look up a name.""" + """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 @@ -186,13 +208,15 @@ def persistent_id(self, obj): 'SafeDisk cannot store a subclass of ndarray: %s.%s' % (type(obj).__module__, type(obj).__qualname__) ) - # Holding `obj` keeps its id from being reused: pickle's own memo - # skips persistent IDs, and a value may hand it temporary arrays. - entry = self._arrays.get(id(obj)) - if entry is None: - entry = (obj, ('npy', _pack_array(numpy, obj))) - self._arrays[id(obj)] = entry - return entry[1] + # 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) @@ -209,7 +233,16 @@ def persistent_id(self, obj): class _ArrayUnpickler(SafeUnpickler): - """Unpickler that reads back what _ArrayPickler wrote.""" + """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) @@ -218,8 +251,8 @@ def __init__(self, *args, **kwargs): def persistent_load(self, pid): # pylint: disable=unidiomatic-typecheck - if type(pid) is tuple: - if len(pid) == 3 and pid[0] == 'scalar': + 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) @@ -230,17 +263,19 @@ def persistent_load(self, pid): entry = (value, _unpack_scalar(pid[1], value)) self._strings[id(value)] = entry return entry[1] - if len(pid) == 2 and pid[0] == 'npy': - payload = pid[1] - if type(payload) is bytearray: + 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 - # `payload` keeps its id from being reused. - entry = self._arrays.get(id(payload)) + # the pair keeps their ids from being reused. + key = (id(header), id(payload)) + entry = self._arrays.get(key) if entry is None: - entry = (payload, _unpack_array(payload)) - self._arrays[id(payload)] = entry - return entry[1] + array = _unpack_array(header, payload) + entry = (header, payload, array) + self._arrays[key] = entry + return entry[2] raise pickle.UnpicklingError('SafeDisk refuses a persistent ID') diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 1a372e8..a461648 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -1078,12 +1078,17 @@ in numpy's ``.npy`` format: cache['embedding'] = numpy.arange(3.0) cache['result'] = {'scores': numpy.zeros(2), 'best': numpy.float64(1.5)} -* An array anywhere in a value is saved with ``allow_pickle=False``, including - inside classes allowed with :meth:`SafeDisk.with_types +* 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``. Its header is checked against the data before it is - read back, also with ``allow_pickle=False``. An array that appears twice in a - value is stored once and is one array when read back. + 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 diff --git a/tests/test_safedisk_numpy.py b/tests/test_safedisk_numpy.py index 022fb83..5084d4e 100644 --- a/tests/test_safedisk_numpy.py +++ b/tests/test_safedisk_numpy.py @@ -13,6 +13,7 @@ import sys import tempfile import tracemalloc +import warnings import pytest @@ -91,10 +92,13 @@ def same(actual, expected): assert actual.tobytes() == expected.tobytes() -def npy_bytes(array): +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) - return buffer.getvalue() + data = buffer.getvalue() + size = len(data) - array.nbytes + return data[:size], data[size:] class Pid: @@ -121,8 +125,9 @@ def persid(pid): return forge(Pid(pid)) -def array_pid(data): - return ('npy', bytearray(data)) +def array_pid(parts): + header, data = parts + return ('npy', header, bytearray(data)) def write_row(directory, data): @@ -160,7 +165,12 @@ def npy_header(header, body=b'', version=1): prefix = 10 if version == 1 else 12 header += b' ' * (64 - (prefix + len(header) + 1) % 64) + b'\n' magic = b'\x93NUMPY' + bytes([version, 0]) - return magic + struct.pack(size, len(header)) + header + body + return magic + struct.pack(size, len(header)) + header, body + + +def raw(header, body=b''): + """Forged parts whose header is not a ``.npy`` header at all.""" + return header, body @pytest.mark.parametrize( @@ -186,6 +196,36 @@ def test_array_round_trip(directory, array): same(cache['key'], array) +def read_only(array): + array.setflags(write=False) + return array + + +@pytest.mark.parametrize( + 'array', + [ + read_only(np.arange(6.0)), + np.arange(12.0)[::2], + np.arange(12.0).reshape(3, 4)[:, 1:3], + np.broadcast_to(np.arange(3.0), (2, 3)), + np.asfortranarray(np.arange(6.0).reshape(2, 3))[:, ::2], + ], + ids=['read-only', 'strided', 'window', 'broadcast', 'fortran-strided'], +) +def test_arrays_not_stored_in_place_are_copied_first(directory, array): + """pickle refuses a non-contiguous buffer and stores a read-only one as + bytes, which the reader refuses, so these take one copy on the way in.""" + with dc.Cache(directory, disk=ArrayDisk) as cache: + cache['key'] = {'a': array, 'b': array} + result = cache['key'] + + assert result['a'] is result['b'] + assert result['a'].flags.writeable + assert result['a'].dtype == array.dtype + assert result['a'].shape == array.shape + assert np.array_equal(result['a'], array) + + def test_arrays_in_containers(directory): array = np.arange(3.0) value = {'a': array, 'b': [array, (array, 1)], 'c': 'text'} @@ -247,6 +287,52 @@ def test_arrays_in_classes_outside_the_allowlist_are_refused(directory): cache['key'] = {'a': np.arange(2.0), 'b': Outside()} +@pytest.mark.parametrize('kind', ['string dtype', 'non-latin-1 field']) +def test_refusing_an_array_emits_no_numpy_warning(directory, kind): + if kind == 'string dtype': + if not hasattr(np, 'dtypes') or not hasattr(np.dtypes, 'StringDType'): + pytest.skip('numpy without StringDType') + array = np.array(['a'], dtype=np.dtypes.StringDType()) + else: + array = np.zeros(1, dtype=[('→', 'f8')]) + + with dc.Cache(directory, disk=ArrayDisk) as cache: + with warnings.catch_warnings(): + warnings.simplefilter('error') + with pytest.raises(TypeError, match='cannot store an array'): + cache['key'] = array + + +def test_only_values_with_reduced_objects_are_scanned(directory, monkeypatch): + """Scanning a pickle costs time in proportion to its size, so a value + whose pickle cannot name anything is stored without one.""" + scanned = [] + check = dc.SafeDisk._check_pickle + + def record(self, data, obj, what): + scanned.append(type(obj)) + return check(self, data, obj, what) + + monkeypatch.setattr(dc.SafeDisk, '_check_pickle', record) + + with dc.Cache(directory, disk=ArrayDisk.with_types(Result)) as cache: + cache['plain'] = {'a': np.arange(2.0), 'b': [np.float64(1.5)]} + assert scanned == [] + cache['reduced'] = Result('name', np.arange(2.0)) + assert scanned == [Result] + + +def test_class_objects_next_to_an_array_are_refused(directory): + """A class pickles as a global without being reduced, so the write + check relies on pickle calling reducer_override for it all the same.""" + with dc.Cache(directory, disk=ArrayDisk) as cache: + with pytest.raises(TypeError): + cache['key'] = [Outside, np.arange(2.0)] + with pytest.raises(TypeError): + cache['key'] = {'f': len, 'a': np.arange(2.0)} + assert len(cache) == 0 + + def test_writing_leaves_the_caller_value_alone(directory): array = np.arange(3.0) value = {'a': array} @@ -697,7 +783,9 @@ def test_empty_array_of_a_dtype_with_no_width(directory): assert cache['key'].shape == (0,) -def test_writing_a_large_array_copies_it_at_most_twice(directory): +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: @@ -706,12 +794,10 @@ def test_writing_a_large_array_copies_it_at_most_twice(directory): _, peak = tracemalloc.get_traced_memory() tracemalloc.stop() - assert peak < array.nbytes * 2.2 + assert peak < array.nbytes * 1.3 -def test_writing_an_array_in_an_allowed_class_copies_it_at_most_twice( - directory, -): +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) @@ -722,7 +808,7 @@ def test_writing_an_array_in_an_allowed_class_copies_it_at_most_twice( tracemalloc.stop() same(cache['key'].vector, array) - assert peak < array.nbytes * 2.2 + assert peak < array.nbytes * 1.3 def test_reading_a_large_array_takes_no_second_copy(directory): @@ -845,7 +931,12 @@ def test_environment_selects_a_with_numpy_disk(monkeypatch): def test_row_format_is_pinned(directory): """Changing what is on disk breaks every cache: this is hand-built.""" array = np.arange(3.0) - write_row(directory, persid(array_pid(npy_bytes(array)))) + header = ( + b"\x93NUMPY\x01\x00v\x00{'descr': ' Date: Mon, 2 Nov 2026 10:14:03 -0800 Subject: [PATCH 10/10] Remove the scanner note and the trusted-directory advice from SECURITY.md and the compatibility contract from CONTRIBUTING.md. --- CONTRIBUTING.md | 14 -------------- SECURITY.md | 4 ---- 2 files changed, 18 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c1365b4..fe131ff 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -36,17 +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 row format `SafeDisk.with_numpy()` writes: a persistent ID of `('npy', , )` or `('scalar', , )` in a pickle of protocol 5 or higher, with `tests/test_safedisk_numpy.py` pinning it; -- 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. Entries a process wrote through `SafeDisk.with_numpy()` that hold numpy data are the one exception to reading them back: 5.6.3 opens the directory and raises `UnpicklingError` for those entries. Golden fixtures written by the 5.6.3 wheel guard this in the test suite. diff --git a/SECURITY.md b/SECURITY.md index b0549da..aa292dc 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -22,10 +22,6 @@ Where this fork stands: - `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.