Coverage for zanj\externals.py: 95%
21 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-14 12:57 -0700
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-14 12:57 -0700
1"""for storing/retrieving an item externally in a ZANJ archive"""
3from __future__ import annotations
5import json
6from typing import IO, Any, Callable, Literal, NamedTuple, get_args
8import numpy as np
9from muutils.json_serialize.json_serialize import ObjectPath
10from muutils.json_serialize.util import JSONitem
12# this is to make type checking work -- it will later be overridden
13_ZANJ_pre = Any
15ZANJ_MAIN: str = "__zanj__.json"
16ZANJ_META: str = "__zanj_meta__.json"
18ExternalItemType = Literal["jsonl", "npy"]
20ExternalItemType_vals = get_args(ExternalItemType)
22ExternalItem = NamedTuple(
23 "ExternalItem",
24 [
25 ("item_type", ExternalItemType),
26 ("data", Any),
27 ("path", ObjectPath),
28 ],
29)
32def load_jsonl(zanj: "LoadedZANJ", fp: IO[bytes]) -> list[JSONitem]: # type: ignore[name-defined] # noqa: F821
33 return [json.loads(line) for line in fp]
36def load_npy(zanj: "LoadedZANJ", fp: IO[bytes]) -> np.ndarray: # type: ignore[name-defined] # noqa: F821
37 return np.load(fp)
40EXTERNAL_LOAD_FUNCS: dict[ExternalItemType, Callable[[_ZANJ_pre, IO[bytes]], Any]] = {
41 "jsonl": load_jsonl,
42 "npy": load_npy,
43}
46def GET_EXTERNAL_LOAD_FUNC(item_type: str) -> Callable[[_ZANJ_pre, IO[bytes]], Any]:
47 if item_type not in EXTERNAL_LOAD_FUNCS:
48 raise ValueError(
49 f"unknown external item type: {item_type}, needs to be one of {EXTERNAL_LOAD_FUNCS.keys()}"
50 )
51 # safe to ignore since we just checked
52 return EXTERNAL_LOAD_FUNCS[item_type] # type: ignore[index]