Coverage for skema/program_analysis/tests/test_comprehension_lambda_cast.py: 100%

1# import json NOTE: json and Path aren't used right now, 

2# from pathlib import Path but will be used in the future 

3from skema.program_analysis.CAST.python.ts2cast import TS2CAST 

4from skema.program_analysis.CAST2FN.model.cast import ( 

5 Assignment, 

6 Attribute, 

7 FunctionDef, 

8 Call, 

9 Var, 

10 Loop, 

11 ModelReturn, 

12 Name, 

13 CASTLiteralValue, 

14 ModelIf, 

15 StructureType, 

16 ScalarType, 

17 Operator 

18) 

19 

20def comp1(): 

21 return """ 

22L = [a for a in range(10)] 

23 """ 

24 

25def comp2(): 

26 return """ 

27D = {a:a*2 for a in range(10)} 

28 """ 

29 

30def lambda1(): 

31 return """ 

32F = lambda x : x * 10 

33 """ 

34 

35def generate_cast(test_file_string): 

36 # use Python to CAST 

37 out_cast = TS2CAST(test_file_string, from_file=False).out_cast 

38 

39 return out_cast 

40 

41def test_comp1(): 

42 comp_cast = generate_cast(comp1()) 

43 

44 # Test basic list comprehension 

45 func_node = comp_cast.nodes[0].body[0] 

46 assert isinstance(func_node, FunctionDef) 

47 assert "%comprehension_list_" in func_node.name.name 

48 assert len(func_node.func_args) == 0 

49 

50 func_node_body = func_node.body 

51 func_asg = func_node.body[0] 

52 func_ret = func_node.body[2] 

53 assert isinstance(func_asg, Assignment) 

54 assert isinstance(func_asg.left, Var) 

55 assert isinstance(func_asg.left.val, Name) 

56 assert "list__temp_" == func_asg.left.val.name 

57 

58 assert isinstance(func_asg.right, CASTLiteralValue) 

59 assert func_asg.right.value_type == StructureType.LIST 

60 assert func_asg.right.value == [] 

61 

62 assert isinstance(func_ret, ModelReturn) 

63 assert isinstance(func_ret.value, Name) 

64 assert "list__temp_" == func_ret.value.name 

65 

66 func_loop = func_node_body[1] 

67 assert isinstance(func_loop, Loop) 

68 assert func_loop.post == [] 

69 

70 # Loop Pre 

71 ############# 

72 func_loop_pre = func_loop.pre 

73 loop_pre_iter = func_loop_pre[0] 

74 assert isinstance(loop_pre_iter, Assignment) 

75 assert isinstance(loop_pre_iter.left, Var) 

76 assert isinstance(loop_pre_iter.left.val, Name) 

77 assert "generated_iter" in loop_pre_iter.left.val.name 

78 

79 assert isinstance(loop_pre_iter.right, Call) 

80 assert isinstance(loop_pre_iter.right.func, Name) 

81 assert loop_pre_iter.right.func.name == "iter" 

82 

83 assert len(loop_pre_iter.right.arguments) == 1 

84 range_call = loop_pre_iter.right.arguments[0] 

85 assert isinstance(range_call, Call) 

86 assert isinstance(range_call.func, Name) 

87 assert range_call.func.name == "range" 

88 assert len(range_call.arguments) == 3 

89 assert isinstance(range_call.arguments[0], CASTLiteralValue) 

90 assert isinstance(range_call.arguments[1], CASTLiteralValue) 

91 assert isinstance(range_call.arguments[2], CASTLiteralValue) 

92 

93 loop_pre_next = func_loop_pre[1] 

94 assert isinstance(loop_pre_next, Assignment) 

95 assert isinstance(loop_pre_next.left, CASTLiteralValue) 

96 assert loop_pre_next.left.value_type == StructureType.TUPLE 

97 assert isinstance(loop_pre_next.left.value[0], Var) 

98 assert isinstance(loop_pre_next.left.value[0].val, Name) 

99 assert loop_pre_next.left.value[0].val.name == "a" 

100 

101 assert isinstance(loop_pre_next.left.value[1], Var) 

102 assert isinstance(loop_pre_next.left.value[1].val, Name) 

103 assert "generated_iter_" in loop_pre_next.left.value[1].val.name 

104 

105 assert isinstance(loop_pre_next.left.value[2], Var) 

106 assert isinstance(loop_pre_next.left.value[2].val, Name) 

107 assert "sc_" in loop_pre_next.left.value[2].val.name 

108 

109 assert isinstance(loop_pre_next.right, Call) 

110 assert isinstance(loop_pre_next.right.func, Name) 

111 assert loop_pre_next.right.func.name == "next" 

112 assert len(loop_pre_next.right.arguments) == 1 

113 assert isinstance(loop_pre_next.right.arguments[0], Var) # NOTE: This should be Name, not Var, but it's ok for now 

114 assert "generated_iter_" in loop_pre_next.right.arguments[0].val.name 

115 

116 

117 # Loop Test 

118 ############# 

119 func_loop_test = func_loop.expr 

120 assert isinstance(func_loop_test, Operator) 

121 assert func_loop_test.op == "ast.Eq" 

122 

123 assert isinstance(func_loop_test.operands[0], Name) 

124 assert "sc_" in func_loop_test.operands[0].name 

125 

126 assert isinstance(func_loop_test.operands[1], CASTLiteralValue) 

127 assert func_loop_test.operands[1].value == False 

128 assert func_loop_test.operands[1].value_type == ScalarType.BOOLEAN 

129 

130 

131 # Loop Body 

132 ############# 

133 func_loop_body = func_loop.body 

134 loop_body_append = func_loop_body[0] 

135 assert isinstance(loop_body_append, Call) 

136 assert isinstance(loop_body_append.func, Attribute) 

137 assert isinstance(loop_body_append.func.attr, Name) 

138 assert loop_body_append.func.attr.name == "append" 

139 

140 assert isinstance(loop_body_append.func.value, Name) 

141 assert "list__temp_" in loop_body_append.func.value.name 

142 

143 assert len(loop_body_append.arguments) == 1 

144 assert isinstance(loop_body_append.arguments[0], Var) 

145 assert loop_body_append.arguments[0].val.name == "a" 

146 

147 loop_body_next = func_loop_body[1] 

148 assert isinstance(loop_body_next, Assignment) 

149 assert isinstance(loop_body_next.left, CASTLiteralValue) 

150 assert loop_body_next.left.value_type == StructureType.TUPLE 

151 assert isinstance(loop_body_next.left.value[0], Var) 

152 assert isinstance(loop_body_next.left.value[0].val, Name) 

153 assert loop_body_next.left.value[0].val.name == "a" 

154 

155 assert isinstance(loop_body_next.left.value[1], Var) 

156 assert isinstance(loop_body_next.left.value[1].val, Name) 

157 assert "generated_iter_" in loop_body_next.left.value[1].val.name 

158 

159 assert isinstance(loop_body_next.left.value[2], Var) 

160 assert isinstance(loop_body_next.left.value[2].val, Name) 

161 assert "sc_" in loop_body_next.left.value[2].val.name 

162 

163 assert isinstance(loop_body_next.right, Call) 

164 assert isinstance(loop_body_next.right.func, Name) 

165 assert loop_body_next.right.func.name == "next" 

166 assert len(loop_body_next.right.arguments) == 1 

167 assert isinstance(loop_body_next.right.arguments[0], Var) # NOTE: This should be Name, not Var, but it's ok for now 

168 assert "generated_iter_" in loop_body_next.right.arguments[0].val.name 

169 

170 call_node = comp_cast.nodes[0].body[1] 

171 assert isinstance(call_node, Assignment) 

172 assert isinstance(call_node.left, Var) 

173 assert isinstance(call_node.left.val, Name) 

174 assert call_node.left.val.name == "L" 

175 

176 assert isinstance(call_node.right, Call) 

177 assert isinstance(call_node.right.func, Name) 

178 assert "%comprehension_list_" in call_node.right.func.name 

179 assert len(call_node.right.arguments) == 0 

180 

181def test_comp2(): 

182 comp_cast = generate_cast(comp2()) 

183 

184 # Test basic dict comprehension 

185 # Very similar to list comp, but using dictionaries instead 

186 func_node = comp_cast.nodes[0].body[0] 

187 assert isinstance(func_node, FunctionDef) 

188 assert "%comprehension_dict_" in func_node.name.name 

189 assert len(func_node.func_args) == 0 

190 

191 func_node_body = func_node.body 

192 func_asg = func_node.body[0] 

193 func_ret = func_node.body[2] 

194 assert isinstance(func_asg, Assignment) 

195 assert isinstance(func_asg.left, Var) 

196 assert isinstance(func_asg.left.val, Name) 

197 assert "dict__temp_" == func_asg.left.val.name 

198 

199 assert isinstance(func_asg.right, CASTLiteralValue) 

200 assert func_asg.right.value_type == StructureType.MAP 

201 assert func_asg.right.value == {} 

202 

203 assert isinstance(func_ret, ModelReturn) 

204 assert isinstance(func_ret.value, Name) 

205 assert "dict__temp_" == func_ret.value.name 

206 

207 func_loop = func_node_body[1] 

208 assert isinstance(func_loop, Loop) 

209 assert func_loop.post == [] 

210 

211 # Loop Pre 

212 ############# 

213 func_loop_pre = func_loop.pre 

214 loop_pre_iter = func_loop_pre[0] 

215 assert isinstance(loop_pre_iter, Assignment) 

216 assert isinstance(loop_pre_iter.left, Var) 

217 assert isinstance(loop_pre_iter.left.val, Name) 

218 assert "generated_iter" in loop_pre_iter.left.val.name 

219 

220 assert isinstance(loop_pre_iter.right, Call) 

221 assert isinstance(loop_pre_iter.right.func, Name) 

222 assert loop_pre_iter.right.func.name == "iter" 

223 

224 assert len(loop_pre_iter.right.arguments) == 1 

225 range_call = loop_pre_iter.right.arguments[0] 

226 assert isinstance(range_call, Call) 

227 assert isinstance(range_call.func, Name) 

228 assert range_call.func.name == "range" 

229 assert len(range_call.arguments) == 3 

230 assert isinstance(range_call.arguments[0], CASTLiteralValue) 

231 assert isinstance(range_call.arguments[1], CASTLiteralValue) 

232 assert isinstance(range_call.arguments[2], CASTLiteralValue) 

233 

234 loop_pre_next = func_loop_pre[1] 

235 assert isinstance(loop_pre_next, Assignment) 

236 assert isinstance(loop_pre_next.left, CASTLiteralValue) 

237 assert loop_pre_next.left.value_type == StructureType.TUPLE 

238 assert isinstance(loop_pre_next.left.value[0], Var) 

239 assert isinstance(loop_pre_next.left.value[0].val, Name) 

240 assert loop_pre_next.left.value[0].val.name == "a" 

241 

242 assert isinstance(loop_pre_next.left.value[1], Var) 

243 assert isinstance(loop_pre_next.left.value[1].val, Name) 

244 assert "generated_iter_" in loop_pre_next.left.value[1].val.name 

245 

246 assert isinstance(loop_pre_next.left.value[2], Var) 

247 assert isinstance(loop_pre_next.left.value[2].val, Name) 

248 assert "sc_" in loop_pre_next.left.value[2].val.name 

249 

250 assert isinstance(loop_pre_next.right, Call) 

251 assert isinstance(loop_pre_next.right.func, Name) 

252 assert loop_pre_next.right.func.name == "next" 

253 assert len(loop_pre_next.right.arguments) == 1 

254 assert isinstance(loop_pre_next.right.arguments[0], Var) # NOTE: This should be Name, not Var, but it's ok for now 

255 assert "generated_iter_" in loop_pre_next.right.arguments[0].val.name 

256 

257 

258 # Loop Test 

259 ############# 

260 func_loop_test = func_loop.expr 

261 assert isinstance(func_loop_test, Operator) 

262 assert func_loop_test.op == "ast.Eq" 

263 

264 assert isinstance(func_loop_test.operands[0], Name) 

265 assert "sc_" in func_loop_test.operands[0].name 

266 

267 assert isinstance(func_loop_test.operands[1], CASTLiteralValue) 

268 assert func_loop_test.operands[1].value == False 

269 assert func_loop_test.operands[1].value_type == ScalarType.BOOLEAN 

270 

271 

272 # Loop Body 

273 ############# 

274 func_loop_body = func_loop.body 

275 loop_body_set = func_loop_body[0] 

276 assert isinstance(loop_body_set, Assignment) 

277 assert isinstance(loop_body_set.left, Var) 

278 assert isinstance(loop_body_set.left.val, Name) 

279 assert "dict__temp_" in loop_body_set.left.val.name 

280 

281 assert isinstance(loop_body_set.right, Call) 

282 assert isinstance(loop_body_set.right.func, Name) 

283 assert loop_body_set.right.func.name == "_set" 

284 assert len(loop_body_set.right.arguments) == 3 

285 

286 assert isinstance(loop_body_set.right.arguments[0], Name) 

287 assert "dict__temp_" in loop_body_set.right.arguments[0].name 

288 

289 assert isinstance(loop_body_set.right.arguments[1], Name) 

290 assert loop_body_set.right.arguments[1].name == "a" 

291 

292 assert isinstance(loop_body_set.right.arguments[2], Operator) 

293 assert loop_body_set.right.arguments[2].op == "ast.Mult" 

294 assert isinstance(loop_body_set.right.arguments[2].operands[0], Name) 

295 assert loop_body_set.right.arguments[2].operands[0].name == "a" 

296 assert isinstance(loop_body_set.right.arguments[2].operands[1], CASTLiteralValue) 

297 assert loop_body_set.right.arguments[2].operands[1].value == '2' 

298 

299 loop_body_next = func_loop_body[1] 

300 assert isinstance(loop_body_next, Assignment) 

301 assert isinstance(loop_body_next.left, CASTLiteralValue) 

302 assert loop_body_next.left.value_type == StructureType.TUPLE 

303 assert isinstance(loop_body_next.left.value[0], Var) 

304 assert isinstance(loop_body_next.left.value[0].val, Name) 

305 assert loop_body_next.left.value[0].val.name == "a" 

306 

307 assert isinstance(loop_body_next.left.value[1], Var) 

308 assert isinstance(loop_body_next.left.value[1].val, Name) 

309 assert "generated_iter_" in loop_body_next.left.value[1].val.name 

310 

311 assert isinstance(loop_body_next.left.value[2], Var) 

312 assert isinstance(loop_body_next.left.value[2].val, Name) 

313 assert "sc_" in loop_body_next.left.value[2].val.name 

314 

315 assert isinstance(loop_body_next.right, Call) 

316 assert isinstance(loop_body_next.right.func, Name) 

317 assert loop_body_next.right.func.name == "next" 

318 assert len(loop_body_next.right.arguments) == 1 

319 assert isinstance(loop_body_next.right.arguments[0], Var) # NOTE: This should be Name, not Var, but it's ok for now 

320 assert "generated_iter_" in loop_body_next.right.arguments[0].val.name 

321 

322 call_node = comp_cast.nodes[0].body[1] 

323 assert isinstance(call_node, Assignment) 

324 assert isinstance(call_node.left, Var) 

325 assert isinstance(call_node.left.val, Name) 

326 assert call_node.left.val.name == "D" 

327 

328 assert isinstance(call_node.right, Call) 

329 assert isinstance(call_node.right.func, Name) 

330 assert "%comprehension_dict_" in call_node.right.func.name 

331 assert len(call_node.right.arguments) == 0 

332 

333 

334def test_lambda1(): 

335 lambda_cast = generate_cast(lambda1()) 

336 

337 # Test basic lambda 

338 func_node = lambda_cast.nodes[0].body[0] 

339 assert isinstance(func_node, FunctionDef) 

340 assert "%lambda_" in func_node.name.name 

341 assert len(func_node.func_args) == 1 

342 assert isinstance(func_node.func_args[0], Var) 

343 assert func_node.func_args[0].val.name == "x" 

344 

345 func_body = func_node.body[0] 

346 assert isinstance(func_body, ModelReturn) 

347 

348 assert isinstance(func_body.value, Operator) 

349 assert isinstance(func_body.value.operands[0], Name) 

350 assert func_body.value.operands[0].name == "x" 

351 

352 assert isinstance(func_body.value.operands[1], CASTLiteralValue) 

353 assert func_body.value.operands[1].value_type == ScalarType.INTEGER 

354 assert func_body.value.operands[1].value == '10' 

355 

356 call_node = lambda_cast.nodes[0].body[1] 

357 assert isinstance(call_node, Assignment) 

358 assert isinstance(call_node.left, Var) 

359 assert isinstance(call_node.left.val, Name) 

360 assert call_node.left.val.name == "F" 

361 

362 assert isinstance(call_node.right, Call) 

363 assert isinstance(call_node.right.func, Name) 

364 assert "%lambda_" in call_node.right.func.name 

365 

366 assert len(call_node.right.arguments) == 1 

367 assert isinstance(call_node.right.arguments[0], Name) 

368 assert call_node.right.arguments[0].name == "x" 

369