| 49 | |
| 50 | |
| 51 | def test_source_ast(): |
| 52 | source = Source(matmul) |
| 53 | mod = source.as_ast() |
| 54 | assert isinstance(mod, doc.Module) |
| 55 | func_def = mod.body[0] |
| 56 | assert isinstance(func_def, doc.FunctionDef) |
| 57 | assert func_def.name == "matmul" |
| 58 | func_args = func_def.args |
| 59 | assert ( |
| 60 | len(func_args.args) == 3 |
| 61 | and func_args.args[0].arg == "a" |
| 62 | and func_args.args[1].arg == "b" |
| 63 | and func_args.args[2].arg == "c" |
| 64 | ) |
| 65 | func_body = func_def.body |
| 66 | assert len(func_body) == 4 |
| 67 | func_assigns = func_body[:3] |
| 68 | assert ( |
| 69 | isinstance(func_assigns[0], doc.Assign) |
| 70 | and func_assigns[0].targets[0].id == "A" |
| 71 | and isinstance(func_assigns[1], doc.Assign) |
| 72 | and func_assigns[1].targets[0].id == "B" |
| 73 | and isinstance(func_assigns[2], doc.Assign) |
| 74 | and func_assigns[2].targets[0].id == "C" |
| 75 | ) |
| 76 | func_for = func_body[3] |
| 77 | assert ( |
| 78 | len(func_for.target.elts) == 3 |
| 79 | and func_for.target.elts[0].id == "i" |
| 80 | and func_for.target.elts[1].id == "j" |
| 81 | and func_for.target.elts[2].id == "k" |
| 82 | ) |
| 83 | for_body = func_for.body |
| 84 | assert len(for_body) == 1 |
| 85 | for_block = for_body[0] |
| 86 | assert isinstance(for_block, doc.With) and len(for_block.body) == 2 |
| 87 | |
| 88 | |
| 89 | def test_nesting_parsing(): |