Coverage for tests\test_zanj_basic.py: 100%

35 statements  

« prev     ^ index     » next       coverage.py v7.6.12, created at 2025-02-14 12:57 -0700

1from __future__ import annotations 

2 

3import json 

4import typing 

5from pathlib import Path 

6 

7import numpy as np 

8import pandas as pd # type: ignore 

9 

10from zanj import ZANJ 

11 

12np.random.seed(0) 

13 

14 

15TEST_DATA_PATH: Path = Path("tests/junk_data") 

16 

17 

18def array_meta(x: typing.Any) -> dict: 

19 if isinstance(x, np.ndarray): 

20 return dict( 

21 shape=list(x.shape), 

22 dtype=str(x.dtype), 

23 contents=str(x), 

24 ) 

25 else: 

26 return dict( 

27 type=type(x).__name__, 

28 contents=str(x), 

29 ) 

30 

31 

32def test_numpy(): 

33 data = dict( 

34 name="testing zanj", 

35 some_array=np.random.rand(128, 128), 

36 some_other_array=np.random.rand(16, 64), 

37 small_array=np.random.rand(4, 4), 

38 ) 

39 fname: Path = TEST_DATA_PATH / "test_numpy.zanj" 

40 z: ZANJ = ZANJ() 

41 z.save(data, fname) 

42 recovered_data = z.read(fname) 

43 

44 print(f"{list(data.keys()) = }") 

45 print(f"{list(recovered_data.keys()) = }") 

46 original_vals: dict = {k: array_meta(v) for k, v in data.items()} 

47 print(json.dumps(original_vals, indent=2)) 

48 recovered_vals: dict = {k: array_meta(v) for k, v in recovered_data.items()} 

49 print(json.dumps(recovered_vals, indent=2)) 

50 

51 assert sorted(list(data.keys())) == sorted(list(recovered_data.keys())) 

52 # assert all([type(data[k]) == type(recovered_data[k]) for k in data.keys()]) 

53 

54 assert all( 

55 [ 

56 data["name"] == recovered_data["name"], 

57 np.allclose(data["some_array"], recovered_data["some_array"]), 

58 np.allclose(data["some_other_array"], recovered_data["some_other_array"]), 

59 np.allclose(data["small_array"], recovered_data["small_array"]), 

60 ] 

61 ), f"assert failed:\n{data = }\n{recovered_data = }" 

62 

63 

64def test_jsonl(): 

65 data = dict( 

66 name="testing zanj jsonl", 

67 iris_data=pd.read_csv("tests/input_data/iris.csv"), 

68 brain_data=pd.read_csv("tests/input_data/brain_networks.csv"), 

69 some_array=np.random.rand(128, 128), 

70 ) 

71 fname: Path = TEST_DATA_PATH / "test_jsonl.zanj" 

72 z: ZANJ = ZANJ() 

73 z.save(data, fname) 

74 recovered_data = z.read(fname) 

75 

76 assert sorted(list(data.keys())) == sorted(list(recovered_data.keys())) 

77 # assert all([type(data[k]) == type(recovered_data[k]) for k in data.keys()]) 

78 

79 assert all( 

80 [ 

81 data["name"] == recovered_data["name"], 

82 np.allclose(data["some_array"], recovered_data["some_array"]), 

83 data["iris_data"].equals(recovered_data["iris_data"]), 

84 data["brain_data"].equals(recovered_data["brain_data"]), 

85 ] 

86 )