Class for recursing down ASTs and printing them in a very simple format, mainly for instructive purposes and, perhaps, debugging.
| 38 | |
| 39 | |
| 40 | class ASTPrinter(ExprFunctor): |
| 41 | """ |
| 42 | Class for recursing down ASTs and printing them in a very simple format, |
| 43 | mainly for instructive purposes and, perhaps, debugging. |
| 44 | """ |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | indent_str=" ", |
| 49 | include_struct_info_annotations=True, |
| 50 | include_call_attrs=True, |
| 51 | ): |
| 52 | self.indent_str = indent_str |
| 53 | self.include_struct_info_annotations = include_struct_info_annotations |
| 54 | self.include_call_attrs = include_call_attrs |
| 55 | |
| 56 | def visit_expr(self, expr: relax.Expr) -> str: |
| 57 | # extend so we also dispatch to bindings and binding blocks, |
| 58 | # a little silly but IRFunctor hasn't been ported to Python |
| 59 | if isinstance(expr, relax.DataflowBlock): |
| 60 | return self.visit_dataflow_block_(expr) |
| 61 | if isinstance(expr, relax.BindingBlock): |
| 62 | return self.visit_binding_block_(expr) |
| 63 | if isinstance(expr, relax.Binding): |
| 64 | return self.visit_binding_(expr) |
| 65 | return super().visit_expr(expr) |
| 66 | |
| 67 | def indent(self, text: str) -> str: |
| 68 | """ |
| 69 | Indent all lines of the input. |
| 70 | """ |
| 71 | if text == "": |
| 72 | return "" |
| 73 | lines = text.split("\n") |
| 74 | return self.indent_str + f"\n{self.indent_str}".join(lines) |
| 75 | |
| 76 | def build_ast_node(self, nodename: str, force_newline=False, **kwargs: str) -> str: |
| 77 | """ |
| 78 | Returns 'nodename(..., fields[i][0]=fields[i][1], ...)' |
| 79 | with appropriate indentation |
| 80 | """ |
| 81 | return self.build_list( |
| 82 | map(lambda field: f"{field[0]}={field[1]}", kwargs.items()), |
| 83 | open_tok=f"{nodename}(", |
| 84 | close_tok=")", |
| 85 | force_newline=force_newline, |
| 86 | ) |
| 87 | |
| 88 | def build_expr(self, node: relax.Expr, nodename: str, force_newline=False, **kwargs: str): |
| 89 | """ |
| 90 | Renders a Relax expression as a string using `build_ast_node`. |
| 91 | Handles whether to include the struct_info fields. |
| 92 | """ |
| 93 | fields = kwargs.copy() |
| 94 | if node.struct_info_ and self.include_struct_info_annotations: |
| 95 | fields["struct_info"] = self.visit_struct_info_(node.struct_info) |
| 96 | return self.build_ast_node(nodename, force_newline=force_newline, **fields) |
| 97 |
no outgoing calls
searching dependent graphs…