| 60 | |
| 61 | |
| 62 | class NL2SQLLanguage(object): |
| 63 | root_type = 'sql' |
| 64 | |
| 65 | def __init__( |
| 66 | self, |
| 67 | asdl_file, |
| 68 | output_from=True, #< changed |
| 69 | use_table_pointer=True, #< changed |
| 70 | include_literals=False, #< changed |
| 71 | include_columns=True, |
| 72 | end_with_from=True, #< changed |
| 73 | clause_order=None, |
| 74 | infer_from_conditions=True, #< changed |
| 75 | factorize_sketch=2): |
| 76 | |
| 77 | # collect pointers and checkers |
| 78 | self.pointers = set(['table', 'column', 'value']) |
| 79 | custom_primitive_type_checkers = {} |
| 80 | custom_primitive_type_checkers['table'] = lambda x: isinstance(x, int) |
| 81 | custom_primitive_type_checkers['column'] = lambda x: isinstance(x, int) |
| 82 | custom_primitive_type_checkers['value'] = lambda x: isinstance(x, int) |
| 83 | |
| 84 | self.include_columns = include_columns |
| 85 | # create ast wrapper |
| 86 | self.factorize_sketch = factorize_sketch |
| 87 | self.ast_wrapper = ast_util.ASTWrapper( |
| 88 | asdl.parse(asdl_file), |
| 89 | custom_primitive_type_checkers=custom_primitive_type_checkers) |
| 90 | |
| 91 | # from field |
| 92 | self.output_from = output_from |
| 93 | self.end_with_from = end_with_from |
| 94 | self.clause_order = clause_order |
| 95 | self.infer_from_conditions = infer_from_conditions |
| 96 | if self.clause_order: |
| 97 | # clause order is prioritized over configurations like end_with_from |
| 98 | assert factorize_sketch == 2 # TODO support other grammars |
| 99 | sql_fields = self.ast_wrapper.product_types['sql'].fields |
| 100 | letter2field = {k: v for k, v in zip("SFWGOI", sql_fields)} |
| 101 | new_sql_fields = [letter2field[k] for k in self.clause_order] |
| 102 | self.ast_wrapper.product_types['sql'].fields = new_sql_fields |
| 103 | else: |
| 104 | if not self.output_from: |
| 105 | sql_fields = self.ast_wrapper.product_types['sql'].fields |
| 106 | assert sql_fields[1].name == 'from' |
| 107 | del sql_fields[1] |
| 108 | else: |
| 109 | sql_fields = self.ast_wrapper.product_types['sql'].fields |
| 110 | assert sql_fields[1].name == "from" |
| 111 | if self.end_with_from: |
| 112 | sql_fields.append(sql_fields[1]) |
| 113 | del sql_fields[1] |
| 114 | |
| 115 | def parse(self, code, section): |
| 116 | return self.parse_sql(code) |
| 117 | |
| 118 | def unparse(self, tree, db, value_list): |
| 119 | unparser = NL2SQLUnparser(self.ast_wrapper, db, value_list, self.factorize_sketch) |