Coverage for tests/unit/no_torch/test_zanj_serializable_dataclass.py: 100%

1 from __future__ import annotations

2

3 import json

4 import sys

5 import typing

6 from pathlib import Path

7

8 import numpy as np

9 import pandas as pd # type: ignore[import]

10 from muutils.json_serialize import (

11 SerializableDataclass,

12 serializable_dataclass,

13 serializable_field,

14)

15

16 from zanj import ZANJ

17

18 np.random.seed(0)

19

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

21

22 SUPPORTS_KW_ONLY: bool = bool(sys.version_info >= (3, 10))

23

24

25@serializable_dataclass

26 class BasicZanj(SerializableDataclass):

27 a: str

28 q: int = 42

29 c: typing.List[int] = serializable_field(default_factory=list)

30

31

32 def test_Basic():

33 instance = BasicZanj("hello", 42, [1, 2, 3])

34

35 z = ZANJ()

36 path = TEST_DATA_PATH / "test_BasicZanj.zanj"

37 z.save(instance, path)

38 recovered = z.read(path)

39 assert instance == recovered

40

41

42@serializable_dataclass

43 class Nested(SerializableDataclass):

44 name: str

45 basic: BasicZanj

46 val: float

47

48

49 def test_Nested():

50 instance = Nested("hello", BasicZanj("hello", 42, [1, 2, 3]), 3.14)

51

52 z = ZANJ()

53 path = TEST_DATA_PATH / "test_Nested.zanj"

54 z.save(instance, path)

55 recovered = z.read(path)

56 assert instance == recovered

57

58

59@serializable_dataclass

60 class Nested_with_container(SerializableDataclass):

61 name: str

62 basic: BasicZanj

63 val: float

64 container: typing.List[Nested] = serializable_field(default_factory=list)

65

66

67 def test_Nested_with_container():

68 instance = Nested_with_container(

69 "hello",

70 basic=BasicZanj("hello", 42, [1, 2, 3]),

71 val=3.14,

72 container=[

73 Nested("n1", BasicZanj("n1_b", 123, [4, 5, 7]), 2.71),

74 Nested("n2", BasicZanj("n2_b", 456, [7, 8, 9]), 6.28),

75 ],

76 )

77

78 z = ZANJ()

79 path = TEST_DATA_PATH / "test_Nested_with_container.zanj"

80 z.save(instance, path)

81 recovered = z.read(path)

82 assert instance == recovered

83

84

85@serializable_dataclass

86 class sdc_with_np_array(SerializableDataclass):

87 name: str

88 arr1: np.ndarray

89 arr2: np.ndarray

90

91

92 def test_sdc_with_np_array_small():

93 instance = sdc_with_np_array("small arrays", np.random.rand(10), np.random.rand(20))

94

95 z = ZANJ()

96 path = TEST_DATA_PATH / "test_sdc_with_np_array.zanj"

97 z.save(instance, path)

98 recovered = z.read(path)

99 assert instance == recovered

100

101

102 def test_sdc_with_np_array():

103 instance = sdc_with_np_array(

104 "bigger arrays", np.random.rand(128, 128), np.random.rand(256, 256)

105 )

106

107 z = ZANJ()

108 path = TEST_DATA_PATH / "test_sdc_with_np_array.zanj"

109 z.save(instance, path)

110 recovered = z.read(path)

111 assert instance == recovered

112

113

114@serializable_dataclass

115 class sdc_with_df(SerializableDataclass):

116 name: str

117 iris_data: pd.DataFrame

118 brain_data: pd.DataFrame

119

120

121 def test_sdc_with_df():

122 instance = sdc_with_df(

123 "downloaded_data",

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

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

126 )

127

128 z = ZANJ()

129 path = TEST_DATA_PATH / "test_sdc_with_df.zanj"

130 z.save(instance, path)

131 recovered = z.read(path)

132 assert instance == recovered

133

134

135@serializable_dataclass

136 class sdc_container_explicit(SerializableDataclass):

137 name: str

138 container: typing.List[Nested] = serializable_field(

139 default_factory=list,

140 # as jsonl string, for whatever reason

141 serialization_fn=lambda c: "\n".join([json.dumps(n.serialize()) for n in c]),

142 loading_fn=lambda data: [

143 Nested.load(json.loads(n)) for n in data["container"].split("\n")

144 ],

145 # TODO: explicitly specifying the following does not work, since it gets automatically converted before we call load in `loading_fn`:

146 # serialization_fn=lambda c: [n.serialize() for n in c],

147 # loading_fn=lambda data: [Nested.load(n) for n in data["container"]],

148 )

149

150

151 def test_sdc_container_explicit():

152 instance = sdc_container_explicit(

153 "container explicit",

154 container=[

155 Nested(

156 f"n-{n}",

157 BasicZanj(f"n-{n}_b", n * 10 + 1, [n + 1, n + 2, n + 10]),

158 n * np.pi,

159 )

160 for n in range(10)

161 ],

162 )

163

164 z = ZANJ()

165 path = TEST_DATA_PATH / "test_sdc_container_explicit.zanj"

166 z.save(instance, path)

167 recovered = z.read(path)

168 assert instance == recovered