Skip to content

feat: try to guess file type rather than trust file ext - #77

Open
McPatate wants to merge 1 commit into
mainfrom
feat/guess_file_type_over_extension
Open

feat: try to guess file type rather than trust file ext#77
McPatate wants to merge 1 commit into
mainfrom
feat/guess_file_type_over_extension

Conversation

@McPatate

Copy link
Copy Markdown
Collaborator

we've been having some amount of errors on the hub scanner side, which this would help mitigate. Rather than having to reimplement in our codebase, imo it makes sense to have this logic here!

Curious to hear your feedback on this. AI-generated with some reviewing/clean up.

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
@McPatate
McPatate requested a review from mmaitre314 July 23, 2026 20:57
@mmaitre314

Copy link
Copy Markdown
Owner

AI-generated PR description for future reference:

Motivation

The hub-side scanner sees a steady trickle of errors from files whose extension
doesn't match their contents — numpy arrays saved as .pkl, .npz archives named
.npy, PyTorch checkpoints with no extension at all. Today picklescan dispatches
almost entirely on extension, so these either hit the wrong parser and raise, or get
scanned as something they aren't.

What this changes

Dispatch on content first, extension second.

  • scan_bytes now peeks the first 8 bytes before doing anything else. \x93NUMPY
    routes to scan_numpy; a pickle PROTO opcode routes to scan_pickle_bytes
    (or scan_pytorch, if the stream carries the legacy torch magic number).
    The extension-based .npy dispatch is dropped — the magic bytes cover it.
  • scan_zip_bytes gives member magic bytes precedence over member extension, so a
    zip entry named data.pkl holding an .npy array is scanned as numpy.
  • New _is_legacy_pytorch helper: non-destructive check for the torch legacy
    serialization magic number, used to detect extensionless PyTorch files. Restores
    stream position on both success and failure.
  • scan_bytes seeks back to the stream's original offset rather than 0, so it no
    longer assumes it was handed a stream positioned at the start.

Two crashes fixed along the way.

  • scan_numpy routes .npz (zip) content to scan_zip_bytes instead of raising
    ValueError. That exception wasn't caught by scan_zip_bytes's handler, so a
    single mislabeled member aborted the whole archive scan.
  • scan_pickle_bytes catches the ValueError that _list_globals raises on a
    STACK_GLOBAL with too few operands on the stack, and reports scan_err instead
    of propagating. This is what non-pickle content fed to the pickle parser looks
    like (\x93 is STACK_GLOBAL), so it became reachable once dispatch got looser.
  • _list_globals no longer performs its rewind seek when the peek hit EOF. On a
    ZipExtFile a backward seek resets the member and re-reads it, and on Python
    3.12+ that reset re-enables the CRC validation RelaxedZipFile deliberately
    disabled — turning a corrupt-CRC member into a BadZipFile after parsing
    succeeded, discarding globals already found.

Testing

Four new tests in tests/test_scanner.py covering numpy content under .pkl/.bin/
no extension, .npz content under .npy, pickle content under .npy, a mislabeled
zip member, extensionless PyTorch (both legacy and zip formats), and the STACK_GLOBAL
underflow case.

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.

2 participants