Coverage for zanj/externals.py: 95%

1"""for storing/retrieving an item externally in a ZANJ archive"""

2

3 from __future__ import annotations

4

5 import json

6 from typing import IO, Any, Callable, Literal, NamedTuple, get_args

7

8 import numpy as np

9 from muutils.json_serialize.json_serialize import ObjectPath

10 from muutils.json_serialize.util import JSONitem

11

12# this is to make type checking work -- it will later be overridden

13 _ZANJ_pre = Any

14

15 ZANJ_MAIN: str = "__zanj__.json"

16 ZANJ_META: str = "__zanj_meta__.json"

17

18 ExternalItemType = Literal["jsonl", "npy"]

19

20 ExternalItemType_vals = get_args(ExternalItemType)

21

22 ExternalItem = NamedTuple(

23 "ExternalItem",

24 [

25 ("item_type", ExternalItemType),

26 ("data", Any),

27 ("path", ObjectPath),

28 ],

29)

30

31

32 def load_jsonl(zanj: "LoadedZANJ", fp: IO[bytes]) -> list[JSONitem]: # type: ignore[name-defined] # noqa: F821

33 return [json.loads(line) for line in fp]

34

35

36 def load_npy(zanj: "LoadedZANJ", fp: IO[bytes]) -> np.ndarray: # type: ignore[name-defined] # noqa: F821

37 return np.load(fp)

38

39

40 EXTERNAL_LOAD_FUNCS: dict[ExternalItemType, Callable[[_ZANJ_pre, IO[bytes]], Any]] = {

41 "jsonl": load_jsonl,

42 "npy": load_npy,

43}

44

45

46 def 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]