From 0f5fd5d55b4d65d06a814842af3228f57c8ae9c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Wed, 22 Jul 2026 13:45:24 +0200 Subject: [PATCH] feat: support caching binary assets via `BlobAsset` --- marimo/_save/loaders/lazy.py | 2 +- marimo/_save/stubs/__init__.py | 7 ++- marimo/_save/stubs/lazy_stub.py | 22 +++++++ tests/_save/stubs/test_lazy_asset_codec.py | 72 ++++++++++++++++++++++ 4 files changed, 101 insertions(+), 2 deletions(-) create mode 100644 tests/_save/stubs/test_lazy_asset_codec.py diff --git a/marimo/_save/loaders/lazy.py b/marimo/_save/loaders/lazy.py index 1765b2a570b..b612e9ed4d1 100644 --- a/marimo/_save/loaders/lazy.py +++ b/marimo/_save/loaders/lazy.py @@ -357,7 +357,7 @@ def to_item( ref = _maybe_import_ref(value) if ref is not None: return Item(import_ref=ref) - if loader in ("pickle", "npy", "arrow", "pt"): + if loader in ("pickle", "npy", "arrow", "pt", "bin"): # Blob strategies: the file extension is the loader name, matching # the path `save_cache` writes (`{var}.{loader}`). Listing them # together keeps a new format (e.g. `pt`) from silently falling diff --git a/marimo/_save/stubs/__init__.py b/marimo/_save/stubs/__init__.py index 6776bba76c5..be0fce1e5d9 100644 --- a/marimo/_save/stubs/__init__.py +++ b/marimo/_save/stubs/__init__.py @@ -7,7 +7,11 @@ from marimo._save.stubs.class_stub import ClassStub from marimo._save.stubs.function_stub import FunctionStub -from marimo._save.stubs.lazy_stub import LAZY_STUB_LOOKUP, ReferenceStub +from marimo._save.stubs.lazy_stub import ( + LAZY_STUB_LOOKUP, + BlobAsset, + ReferenceStub, +) from marimo._save.stubs.module_stub import ModuleStub from marimo._save.stubs.pydantic_stub import PydanticStub from marimo._save.stubs.stubs import ( @@ -78,6 +82,7 @@ def maybe_get_custom_stub(value: Any) -> CustomStub | None: __all__ = [ "CUSTOM_STUBS", "LAZY_STUB_LOOKUP", + "BlobAsset", "ClassStub", "CustomStub", "FunctionStub", diff --git a/marimo/_save/stubs/lazy_stub.py b/marimo/_save/stubs/lazy_stub.py index 0305090f5b1..e6eb71abf57 100644 --- a/marimo/_save/stubs/lazy_stub.py +++ b/marimo/_save/stubs/lazy_stub.py @@ -119,6 +119,13 @@ class Cache(msgspec.Struct, omit_defaults=True): ui_defs: list[str] = msgspec.field(default_factory=list) +class BlobAsset(msgspec.Struct, frozen=True): + data: bytes + media_type: str | None = None + filename: str | None = None + metadata: dict[str, Any] = msgspec.field(default_factory=dict) + + # --------------------------------------------------------------------------- # Unified type-to-loader registry # --------------------------------------------------------------------------- @@ -135,6 +142,7 @@ class Cache(msgspec.Struct, omit_defaults=True): # "npy" — numpy .npy blob # "arrow" — Arrow IPC .arrow blob (polars and pandas) # "pt" — torch .pt blob (torch.save / torch.load) +# "bin" — media-typed bytes with export metadata LAZY_STUB_LOOKUP: dict[str, str] = { "builtins.int": "inline", "builtins.str": "inline", @@ -160,6 +168,7 @@ class Cache(msgspec.Struct, omit_defaults=True): # Subclasses (e.g. torch.nn.Parameter) resolve here through the MRO # walk in maybe_update_lazy_stub; torch.save round-trips them intact. "torch.Tensor": "pt", + "marimo._save.stubs.lazy_stub.BlobAsset": "bin", } @@ -213,11 +222,17 @@ def _pickle_load(data: bytes, type_hint: str | None = None) -> Any: return pickle.loads(data) +def _bin_load(data: bytes, type_hint: str | None = None) -> BlobAsset: + del type_hint + return msgspec.msgpack.decode(data, type=BlobAsset) + + BLOB_DESERIALIZERS: dict[str, Callable[[bytes, str | None], Any]] = { ".pickle": _pickle_load, ".npy": _npy_load, ".arrow": _arrow_load, ".pt": _pt_load, + ".bin": _bin_load, } # --------------------------------------------------------------------------- @@ -282,11 +297,18 @@ def _pt_dump(obj: Any) -> bytes: return buf.getvalue() +def _bin_dump(obj: Any) -> bytes: + if not isinstance(obj, BlobAsset): + raise TypeError(f"expected BlobAsset, got {type(obj).__name__}") + return msgspec.msgpack.encode(obj) + + BLOB_SERIALIZERS: dict[str, Callable[[Any], bytes]] = { "pickle": pickle.dumps, "npy": _npy_dump, "arrow": _arrow_dump, "pt": _pt_dump, + "bin": _bin_dump, } # --------------------------------------------------------------------------- diff --git a/tests/_save/stubs/test_lazy_asset_codec.py b/tests/_save/stubs/test_lazy_asset_codec.py new file mode 100644 index 00000000000..3fbe17fe93a --- /dev/null +++ b/tests/_save/stubs/test_lazy_asset_codec.py @@ -0,0 +1,72 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +from pathlib import Path + +from marimo._save.cache import MARIMO_CACHE_VERSION, Cache +from marimo._save.hash import HashKey +from marimo._save.loaders import LazyLoader +from marimo._save.loaders.lazy import maybe_update_lazy_stub +from marimo._save.stubs.lazy_stub import ( + BLOB_DESERIALIZERS, + BLOB_SERIALIZERS, + BlobAsset, +) +from tests._save.store.mocks import MockStore + + +def test_blob_asset_resolves_to_bin_codec() -> None: + asset = BlobAsset( + data=b"", + media_type="image/svg+xml", + filename="chart.svg", + ) + + assert maybe_update_lazy_stub(asset) == "bin" + + +def test_blob_asset_codec_round_trip() -> None: + asset = BlobAsset( + data=b'{"ok": true}', + media_type="application/json", + filename="data.json", + metadata={"format_id": "example.json.v1"}, + ) + + encoded = BLOB_SERIALIZERS["bin"](asset) + + assert BLOB_DESERIALIZERS[".bin"](encoded, None) == asset + + +def test_lazy_loader_round_trips_blob_asset() -> None: + store = MockStore() + loader = LazyLoader("asset-test", store=store, mode="off") + asset = BlobAsset( + data=b"\x89PNG\r\n\x1a\n", + media_type="image/png", + filename="preview.png", + metadata={"format_id": "image.png.v1"}, + ) + cache = Cache( + defs={"preview": asset}, + hash="asset_hash", + cache_type="Pure", + stateful_refs=set(), + hit=False, + meta={"version": MARIMO_CACHE_VERSION, "return": asset}, + ) + + assert loader.save_cache(cache) + loader.flush() + + assert ( + Path("asset-test") / "asset_hash" / "preview.bin" + ).as_posix() in store._cache + assert ( + Path("asset-test") / "asset_hash" / "return.bin" + ).as_posix() in store._cache + loaded = loader.load_cache(HashKey("asset_hash", "Pure")) + + assert loaded is not None + assert loaded.defs == {"preview": asset} + assert loaded.meta["return"] == asset