Coverage for tests/unit/json_serialize/test_array.py: 100%

1import numpy as np 

2import pytest 

3 

4from muutils.json_serialize import JsonSerializer 

5from muutils.json_serialize.array import ( 

6 ArrayMode, 

7 arr_metadata, 

8 array_n_elements, 

9 load_array, 

10 serialize_array, 

11) 

12 

13# pylint: disable=missing-class-docstring 

14 

15 

16class TestYourModule: 

17 def setup_method(self): 

18 self.array_1d = np.array([1, 2, 3]) 

19 self.array_2d = np.array([[1, 2], [3, 4]]) 

20 self.array_3d = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=np.int64) 

21 self.array_zero_dim = np.array(42) 

22 self.jser = JsonSerializer(array_mode="list") 

23 

24 def test_array_n_elements(self): 

25 assert array_n_elements(self.array_1d) == 3 

26 assert array_n_elements(self.array_2d) == 4 

27 assert array_n_elements(self.array_3d) == 8 

28 assert array_n_elements(self.array_zero_dim) == 1 

29 

30 def test_arr_metadata(self): 

31 metadata = arr_metadata(self.array_3d) 

32 assert metadata["shape"] == [2, 2, 2] 

33 assert metadata["dtype"] == "int64" 

34 assert metadata["n_elements"] == 8 

35 

36 @pytest.mark.parametrize( 

37 "array_mode,expected_type", 

38 [ 

39 ("list", list), 

40 ("array_list_meta", dict), 

41 ("array_hex_meta", dict), 

42 ("array_b64_meta", dict), 

43 ], 

44 ) 

45 def test_serialize_array(self, array_mode: ArrayMode, expected_type: type): 

46 result = serialize_array( 

47 self.jser, self.array_2d, "test_path", array_mode=array_mode 

48 ) 

49 assert isinstance(result, expected_type) 

50 

51 def test_load_array(self): 

52 serialized_array = serialize_array( 

53 self.jser, self.array_3d, "test_path", array_mode="array_list_meta" 

54 ) 

55 loaded_array = load_array(serialized_array, array_mode="array_list_meta") 

56 assert np.array_equal(loaded_array, self.array_3d) 

57 

58 def test_serialize_load_integration(self): 

59 for array_mode in [ 

60 "list", 

61 "array_list_meta", 

62 "array_hex_meta", 

63 "array_b64_meta", 

64 ]: 

65 for array in [self.array_1d, self.array_2d, self.array_3d]: 

66 serialized_array = serialize_array( 

67 self.jser, 

68 array, 

69 "test_path", 

70 array_mode=array_mode, # type: ignore[arg-type] 

71 ) 

72 loaded_array = load_array(serialized_array, array_mode=array_mode) # type: ignore[arg-type] 

73 assert np.array_equal(loaded_array, array) 

74 

75 def test_serialize_load_zero_dim(self): 

76 for array_mode in [ 

77 "list", 

78 "array_list_meta", 

79 "array_hex_meta", 

80 "array_b64_meta", 

81 ]: 

82 serialized_array = serialize_array( 

83 self.jser, 

84 self.array_zero_dim, 

85 "test_path", 

86 array_mode=array_mode, # type: ignore[arg-type] 

87 ) 

88 loaded_array = load_array(serialized_array) 

89 assert np.array_equal(loaded_array, self.array_zero_dim)