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

11 from zanj.consts import JSONitem

12

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

14 _ZANJ_pre = Any

15

16 ZANJ_MAIN: str = "__zanj__.json"

17 ZANJ_META: str = "__zanj_meta__.json"

18

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

20

21 ExternalItemType_vals = get_args(ExternalItemType)

22

23 ExternalItem = NamedTuple(

24 "ExternalItem",

25 [

26 ("item_type", ExternalItemType),

27 ("data", Any),

28 ("path", ObjectPath),

29 ],

30)

31

32

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

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

35

36

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

38 return np.load(fp)

39

40

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

42 "jsonl": load_jsonl,

43 "npy": load_npy,

44}

45

46

47 def GET_EXTERNAL_LOAD_FUNC(item_type: str) -> Callable[[_ZANJ_pre, IO[bytes]], Any]:

48 if item_type not in EXTERNAL_LOAD_FUNCS:

49 raise ValueError(

50 f"unknown external item type: {item_type}, needs to be one of {EXTERNAL_LOAD_FUNCS.keys()}"

51 )

52 # safe to ignore since we just checked

53 return EXTERNAL_LOAD_FUNCS[item_type] # type: ignore[index]