Coverage for tests/assertions.py: 70%

48 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-02-28 12:51 +1100

1""" 

2This module containts various complex assertion routines which are used in unit tests 

3""" 

4 

5import numpy as np 

6import xarray as xr 

7 

8 

9# pylint: disable=too-many-arguments 

10def assert_dataarray_equal( 

11 left, 

12 right, 

13 check_name=True, 

14 check_attrs_values=True, 

15 check_attrs_order=False, 

16 check_dtype=True, 

17 decimals=None, 

18): 

19 """ 

20 Check that two Xarray DataArrays are equal. 

21 

22 Args: 

23 left (xarray.DataArray): a DataArray 

24 right (xarray.DataArray): another DataArray 

25 check_name (Optional[bool]): whether to check the name property of 

26 the DataArrays, defaults to True 

27 check_attrs_values (Optional[bool]): whether to check DataArray 

28 attributes, defaults to True 

29 check_attrs_order (Optional[bool]): whether to check the order of the 

30 DataArray attributes. The order can be checked without checking 

31 the values. Defaults to False 

32 check_dtype(Optional[bool]): If True (default), then checks whether 

33 DataArrays have the same dtype. 

34 decimals (Optional[int]): If supplied, then the DataArrays are rounded 

35 to this many decimal places when testing equality (see np.round). 

36 

37 Returns: 

38 None 

39 """ 

40 if not isinstance(left, xr.DataArray): 40 ↛ 41line 40 didn't jump to line 41, because the condition on line 40 was never true

41 raise TypeError(f"left must be an xarray.DataArray, not {type(left)}") 

42 if not isinstance(right, xr.DataArray): 42 ↛ 43line 42 didn't jump to line 43, because the condition on line 42 was never true

43 raise TypeError(f"right must be an xarray.DataArray, not {type(right)}") 

44 

45 # check that the types of left and right are compatible 

46 assert isinstance(left, type(right)) or isinstance( 

47 right, type(left) 

48 ), f"Incompatible types: left: {type(left)}, right: {type(right)}" 

49 

50 # remember the object type 

51 prefix = type(left).__name__ 

52 

53 # if decimals are supplied, do a rounding, otherwise rounding is just a dummy 'identity' func 

54 # pylint: disable=unnecessary-lambda-assignment 

55 rounding = lambda x: x if decimals is None else np.round(x, decimals=decimals) 

56 

57 decimal_str = "" if decimals is None else f" to {decimals} decimal places" 

58 

59 # check the values using xarray.DataArray.equals or xarray.Dataset.equals 

60 assert rounding(left).equals( 

61 rounding(right) 

62 ), f"{prefix}s are not equal{decimal_str}: \nleft: {left}\nright: {right}\n" 

63 

64 # check the Dataset or DataArray attributes 

65 if check_attrs_values: 65 ↛ 73line 65 didn't jump to line 73, because the condition on line 65 was never false

66 np.testing.assert_equal( 

67 left.attrs, 

68 right.attrs, 

69 err_msg=f"{prefix} attributes are not equal", 

70 ) 

71 

72 # check the attributes order 

73 if check_attrs_order: 73 ↛ 74line 73 didn't jump to line 74, because the condition on line 73 was never true

74 left_keys = list(left.attrs.keys()) 

75 right_keys = list(right.attrs.keys()) 

76 assert left_keys == right_keys, ( 

77 f"order of {prefix} attributes are different:\n" f"left: {left_keys}\n" f"right: {right_keys}\n" 

78 ) 

79 

80 # check that the names of the dataarrays are equal 

81 if check_name: 81 ↛ 84line 81 didn't jump to line 84, because the condition on line 81 was never false

82 assert left.name == right.name, f"DataArray names are not equal:\nleft: {left.name}\nright: {right.name}\n" 

83 

84 if check_dtype: 84 ↛ exitline 84 didn't return from function 'assert_dataarray_equal', because the condition on line 84 was never false

85 assert left.dtype == right.dtype, ( 

86 f"DataArray dtypes are not equal:\nleft: {left.dtype}\nright: " f"{right.dtype}\n" 

87 ) 

88 

89 

90# pylint: disable=too-many-arguments 

91def assert_dataset_equal( 

92 left, 

93 right, 

94 check_ds_attrs_values=True, 

95 check_da_attrs_values=True, 

96 check_attrs_order=False, 

97 check_dtype=True, 

98 decimals=None, 

99): 

100 """ 

101 Assert that two Xarray datasets are equal 

102 

103 Args: 

104 left (xarray.Dataset): a Dataset 

105 right (xarray.Dataset): another Dataset 

106 check_ds_attrs_values (Optional[bool]): whether to check the Dataset 

107 attributes, defaults to True 

108 check_da_attrs_values (Optional[bool]): whether to check the DataArray 

109 attributes of each data variable in the Datasets, defaults to True 

110 check_attrs_order (Optional[bool]): whether to check the order of the 

111 Dataset and/or DataArray attributes. The order can be checked 

112 without checking the values. Defaults to False 

113 check_dtype(Optional[bool]): If True (default), then checks whether 

114 the data variables have the same dtype. 

115 decimals (Optional[int]): If supplied, then the data variables are 

116 rounded to this many decimal places when testing equality 

117 (see np.round). 

118 

119 Returns: 

120 None 

121 """ 

122 if not isinstance(left, xr.Dataset): 122 ↛ 123line 122 didn't jump to line 123, because the condition on line 122 was never true

123 raise TypeError(f"left must be an xarray.Dataset, not {type(left)}") 

124 if not isinstance(right, xr.Dataset): 124 ↛ 125line 124 didn't jump to line 125, because the condition on line 124 was never true

125 raise TypeError(f"right must be an xarray.Dataset, not {type(right)}") 

126 

127 _assert_xarray_equal( 

128 left, 

129 right, 

130 check_attrs_values=check_ds_attrs_values, 

131 check_attrs_order=check_attrs_order, 

132 decimals=decimals, 

133 ) 

134 

135 if check_da_attrs_values or check_attrs_order or check_dtype: 135 ↛ exitline 135 didn't return from function 'assert_dataset_equal', because the condition on line 135 was never false

136 for da_name in left.data_vars: 

137 da_left = left[da_name] 

138 da_right = right[da_name] 

139 try: 

140 assert_dataarray_equal( 

141 da_left, 

142 da_right, 

143 check_attrs_values=check_da_attrs_values, 

144 check_attrs_order=check_attrs_order, 

145 check_dtype=check_dtype, 

146 decimals=decimals, 

147 ) 

148 except AssertionError as exc: 

149 raise AssertionError(f'Dataset variables "{da_name}" are not equal:\n{str(exc)}') from exc 

150 

151 

152# pylint: disable=unnecessary-lambda-assignment 

153def _assert_xarray_equal(left, right, check_attrs_values=True, check_attrs_order=False, decimals=None): 

154 """ 

155 Check that two Xarray objects (Dataset or DataArray) are equal. 

156 

157 Args: 

158 left (xarray.DataArray or xarray.Dataset): left object 

159 right (xarray.DataArray or xarray.Dataset): right object 

160 check_attrs_values (Optional[bool]): whether to check `OrderedDict`s 

161 `left.attrs` against `right.attrs`. Defaults to True 

162 check_attrs_order (Optional[bool]): whether to check the order of keys 

163 in the `OrderedDict`s `left.attrs` and `right.attrs`. The order 

164 can be checked without checking the values. Defaults to False 

165 decimals (Optional[int]): If supplied, then the data objects are 

166 rounded to this many decimal places (see np.round). 

167 

168 Returns: 

169 None 

170 """ 

171 

172 # check that the types of left and right are compatible 

173 assert isinstance(left, type(right)) or isinstance( 

174 right, type(left) 

175 ), f"Incompatible types: left: {type(left)}, right: {type(right)}" 

176 

177 # remember the object type 

178 prefix = type(left).__name__ 

179 

180 # if decimals are supplied, do a rounding, otherwise rounding is just a dummy 'identity' func 

181 rounding = lambda x: x if decimals is None else np.round(x, decimals=decimals) 

182 

183 decimal_str = "" if decimals is None else f" to {decimals} decimal places" 

184 

185 # check the values using xarray.DataArray.equals or xarray.Dataset.equals 

186 assert rounding(left).equals( 

187 rounding(right) 

188 ), f"{prefix}s are not equal{decimal_str}: \nleft: {left}\nright: {right}\n" 

189 

190 # check the Dataset or DataArray attributes 

191 if check_attrs_values: 191 ↛ 199line 191 didn't jump to line 199, because the condition on line 191 was never false

192 np.testing.assert_equal( 

193 left.attrs, 

194 right.attrs, 

195 err_msg=f"{prefix} attributes are not equal", 

196 ) 

197 

198 # check the attributes order 

199 if check_attrs_order: 199 ↛ 200line 199 didn't jump to line 200, because the condition on line 199 was never true

200 left_keys = list(left.attrs.keys()) 

201 right_keys = list(right.attrs.keys()) 

202 assert left_keys == right_keys, ( 

203 f"order of {prefix} attributes are different:\n" f"left: {left_keys}\n" f"right: {right_keys}\n" 

204 )