Fix lineno issues (apache#55)

cyx-6 · Jul 13, 2022 · 9af443d · 9af443d
1 parent 9f42c5d
commit 9af443d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/python/tvm/script/parse/diagnostics.py b/python/tvm/script/parse/diagnostics.py
@@ -40,8 +40,8 @@ def _emit(self, node: doc.AST, message: str, level: diagnostics.DiagnosticLevel)
                     source_name=SourceName(self.source.source_name),
                     line=node.lineno,
                     end_line=node.end_lineno,
-                    column=node.col_offset,
-                    end_column=node.end_col_offset,
+                    column=node.col_offset + 1,
+                    end_column=node.end_col_offset + 1,
                 ),
                 message=message,
             )

diff --git a/python/tvm/script/parse/evaluator.py b/python/tvm/script/parse/evaluator.py
@@ -58,14 +58,26 @@ def _visit(self, node: doc.AST) -> Any:
             if node.id not in self.value_table:
                 self.parser.report_error(node, "Undefined variable: %s" % node.id)
             return node
-        if isinstance(node, doc.Constant):
+        if (not isinstance(node, doc.expr)) or isinstance(
+            node,
+            (
+                doc.Constant,
+                doc.expr_context,
+                doc.operator,
+                doc.boolop,
+                doc.unaryop,
+                doc.cmpop,
+                doc.slice,
+            ),
+        ):
             return node
         new_fields = {}
         for field in node.__class__._FIELDS:  # pylint: disable=protected-access
-            if isinstance(field, doc.AST):
-                new_fields[field] = self._visit(getattr(node, field))
+            attr = getattr(node, field)
+            if isinstance(attr, (doc.AST, tuple, list)):
+                new_fields[field] = self._visit(attr)
             else:
-                new_fields[field] = getattr(node, field)
+                new_fields[field] = attr
         try:
             new_value = _eval_expr(node.__class__(**new_fields), self.value_table)
         except Exception as e: