Skip to content

Add SafeDisk.with_numpy() to store numpy arrays without pickling them - #34

Merged
avgvi merged 10 commits into
masterfrom
feat/safe-array-disk
Nov 2, 2026
Merged

avgvi merged 10 commits into
masterfrom
feat/safe-array-disk

Conversation

@avgvi

@avgvi avgvi commented Nov 2, 2026 •

Copy link
Copy Markdown
Member

SafeDisk refuses numpy, because unpickling an array is not memory-safe, so a library that caches arrays can't use it. Today each one has to convert arrays to bytes by hand, and numpy scalars such as numpy.float64, which pandas results are full of, raise TypeError at write time.

SafeDisk.with_numpy() returns a SafeDisk subclass that stores arrays without pickling them. It follows with_types and with_options, so the three combine in any order: SafeDisk.with_numpy().with_types(Result).with_options(on_error='miss') reads as "safe pickle, plus numpy, plus my class, and treat errors as misses". The pickler and unpickler live in a new diskcache_community/arrays.py. SafeDisk itself still refuses numpy, UNSAFE_MODULES stays, and the with_types error now points to with_numpy().

  • Arrays are written through pickle's persistent IDs as their .npy header and their data, anywhere in a value: inside any container SafeDisk accepts, and inside classes allowed with with_types, such as a dataclass holding arrays.
  • numpy booleans, integers, floats, complex numbers and strings read back as the same numpy type with the same value, including NaN, so moving from Disk doesn't change what code gets back.
  • An array used twice in a value is stored once and is one array when read back. Cycles work.
  • Writing a large array copies it once, into the row, because pickle stores the array's own buffer in place, and reading takes no second copy: the row is read from its file and the array is a writeable view of the stored buffer. Peak memory and time match the default Disk both ways. An array that is not contiguous or not writeable is copied once more first.
  • Object arrays, ndarray subclasses, datetime64, timedelta64 and longdouble scalars, and arrays the reader would refuse, such as a dtype with no width or a header over numpy's limit, raise TypeError at write time. numpy.save would silently drop a mask, so subclasses are refused rather than stored.
  • SafeJSONDisk.with_numpy() raises TypeError, since JSON holds no arrays.
  • numpy stays optional. A value with no numpy in it is an ordinary pickle and costs nothing extra.

This also fixes a bug on every disk: a float NaN stored as a value read back as None, because SQLite stores NaN as NULL. It is now stored as a pickle, which 5.6.3 reads too.

For existing SafeDisk users

Nothing changes unless with_numpy() is called. A plain SafeDisk writes the same rows as before and refuses numpy as before, with the same TypeError. The three new names on the class are with_numpy, the safe_numpy attribute, which is False, and the private _plain_unpickler. The one visible change for everyone is the NaN fix above, where None becomes nan. A directory written by SafeDisk.with_numpy() is read by a plain SafeDisk except for the entries holding numpy data, which raise UnpicklingError or read as a miss with on_error='miss'.

On-disk format

This is the one part that can't change after release. A row holding numpy data is a pickle of protocol 5 or higher containing persistent IDs of ('npy', <bytes of the .npy header>, <bytearray of the array data>) or ('scalar', <type name>, <value>). Keeping the header apart from the data is what lets pickle write the array's memory in place. CONTRIBUTING.md now lists this in the compatibility contract, and test_row_format_is_pinned builds such rows by hand. 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.

An entry holding numpy data can't be read by another disk or by 5.6.3. They open the directory and raise UnpicklingError for that entry, and a disk with on_error='miss' deletes it. Every process that shares a directory needs this disk, and the tutorial says so. Plain entries in the same directory stay readable by everything. Arrays pickled earlier by the default Disk are refused like any other pickle outside the allowlist.

Reading forged entries

Reading is a new path for bytes that may have been tampered with:

  • Sizes in the .npy header are checked against the data before anything is built from it, object dtypes are refused there, a zero-width dtype can't claim a shape, and a header with bytes after it is refused.
  • A forged row can follow a stored array with BUILD, which calls ndarray.__setstate__, the call that makes numpy pickles memory-unsafe. It fails with TypeError because that call accepts only a numpy.dtype in its dtype slot and nothing the reader can build is one: not an allowlisted name, not a persistent ID. A test pins this with five forgeries, and a comment on the unpickler says why it holds.
  • Every forged entry raises UnpicklingError, or reads as a miss with on_error='miss'. That includes a crafted header that makes the Python parser raise MemoryError.
  • Each array buffer and each scalar string decodes once, so a row can't decode to more than the data it holds, and shared structure stays linear because pickle's own memo does the work. A scalar's value must be of the Python type its name is built from, so ('scalar', 'bytes', n) can't ask numpy for n zero bytes.
  • Keys never hold numpy, so key rows use the plain SafeDisk reader and a persistent ID there is refused.
  • A process that can't import numpy raises ImportError even with on_error='miss' and keeps the entry, while a pickle naming a removed class still reads as a miss as documented.

Also in this PR

  • core.py gets small hooks: Disk._store_pickle and SafeDisk._check_pickle, SafeDisk._unpickler now takes a file so file-backed rows are streamed, and _ReaderMissing. core.py imports arrays.py at its end, since arrays.py imports the classes above.
  • The write-time scan that checks the names a pickle looks up now reads past strings and buffers over 4 KB instead of copying them, so an array inside a with_types class peaks at the same 2.1x as a bare one rather than 4x. This applies to every SafeDisk.with_types write. The pylint module limit for core.py rises from 4000 to 4100 lines.
  • DISKCACHE_SAFE_TYPES and DISKCACHE_SAFE_MODULES now apply when DISKCACHE_DISK names any SafeDisk subclass, so a class built with with_numpy() in your own module, such as myapp.cache:ArrayDisk, can be configured from the environment. There is no new environment variable.
  • The docs replace the hand-written numpy.save recipe, SECURITY.md and CONTRIBUTING.md are updated, and CHANGELOG.md has entries under Unreleased.
  • CI installs numpy in the Django and SQLite jobs, so these tests run there and count towards coverage.

Limits

  • numpy values in a key raise TypeError, including memoize arguments.
  • cache.incr doesn't work on a numpy integer stored as a value, as on the default Disk.
  • Writing a 100 MB array matches the default Disk: 14 ms and +96 MB peak RSS against 17 ms and +96 MB, each in a fresh process. Reading matches too: 6.0 ms against 5.8 ms, with no copy beyond the row.
  • A refused array emits no numpy warning before its TypeError.
  • Each array costs about 9 µs to read and 12 µs to write, almost all of it numpy parsing the .npy header, so many small arrays are slower than pickling: a dict of eight 100-element arrays reads in 78 µs against 7 µs on the default Disk, and 10,000 eight-element arrays in one value write in 127 ms and read in 90 ms, against 7 ms and 11 ms.
  • Reading 1M numpy.float64 in a list takes 186 ms, against 70 ms. Writing is about the same, 667 ms against 630 ms.
  • Values with no numpy in them read within 4% of SafeDisk and write within noise.

Checked

  • The full test suite passes on Python 3.10, 3.12 and 3.14 on macOS (805 passed, 2 skipped, 99.1% coverage, arrays.py at 100%), with numpy 2.2 and 2.5.
  • tests/test_safedisk_numpy.py also passes on numpy 1.26 on Python 3.12, and on numpy 1.24 on Python 3.8.
  • 32 forged array entries and 39 forged persistent IDs are each refused with UnpicklingError, and read as a clean miss with on_error='miss'. Read-only, strided, windowed, broadcast and Fortran-strided arrays round-trip. 1001 scalar IDs naming one 100 KB string decode to one numpy string, and a 1.9 GB decode of a 9.5 MB row before this check is now 10 MB.
  • About 8.8 million mutated .npy inputs across four Python and numpy pairs were either accepted or raised UnpicklingError, with no crash, and accepted arrays were read back in full.
  • About 538,000 mutated rows through the array unpickler, on numpy 1.26 and 2.5, gave no crash and no hang. 586 raised MemoryError. A sample of 80 of those did the same with arrays disabled, so the ones checked don't come from this code.
  • A five-round interleaved benchmark against master, with an A/A pair as the noise floor, shows no change on the default Disk or on SafeDisk across the classic, sizes, hit-rate, process, thread and Django workloads, apart from the NaN fix costing about 0.3 µs per NaN read.
  • A directory written by this disk opens on diskcache 5.6.3, which reads the plain entries and a float NaN, and raises UnpicklingError for the array entry.
  • Removing each of 58 protections in turn makes tests fail.
  • bluecheck, doc8, docs, flake8, isortcheck, mypy and pylint pass, and so do the coexistence and golden tests with upstream diskcache 5.6.3 installed alongside.

@read-the-docs-community

read-the-docs-community Bot commented Nov 2, 2026 •

Copy link
Copy Markdown

Documentation build overview

📚 diskcache-community | 🛠️ Build #35589912 | 📁 Comparing 678edce against latest (d8335f2)

  🔍 Preview build  

4 files changed
± api.html
± fix-cve-2025-69872.html
± genindex.html
± tutorial.html

@avgvi avgvi changed the title Add SafeArrayDisk, a SafeDisk that stores numpy arrays as .npy bytes Add SafeDisk.with_numpy() to store numpy arrays without pickling them Nov 2, 2026
avgvi added 3 commits November 2, 2026 08:30
…type name, and read past large strings and buffers in the write-time pickle scan.
…iting copies it once, and refuse BUILD on a stored array by test.
…Y.md and the compatibility contract from CONTRIBUTING.md.
@avgvi
avgvi merged commit 1e463df into master Nov 2, 2026
23 checks passed
@avgvi
avgvi deleted the feat/safe-array-disk branch November 2, 2026 18:22
@avgvi avgvi mentioned this pull request Nov 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant