306 lines
9.9 KiB
Python
306 lines
9.9 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E741
|
|
"""TVM Script Parser Source and diagnostics"""
|
|
|
|
import inspect
|
|
import sys
|
|
|
|
from tvm.error import DiagnosticError
|
|
|
|
from . import doc
|
|
|
|
|
|
class Source:
|
|
"""Source code class for TVMScript.
|
|
|
|
It is constructed by source code str or doc AST tree.
|
|
|
|
Parameters
|
|
----------
|
|
source_name : str
|
|
The filename of the file where the source code locates.
|
|
|
|
start_line : int
|
|
The first line number of the source code.
|
|
|
|
start_column : int
|
|
The first column number of the first line of the source code.
|
|
|
|
source : str
|
|
The source code str of source code.
|
|
|
|
full_source : str
|
|
The complete source code of the file where the source code locates.
|
|
"""
|
|
|
|
source_name: str
|
|
start_line: int
|
|
start_column: int
|
|
source: str
|
|
full_source: str
|
|
|
|
def __init__(self, program: str | doc.AST):
|
|
if isinstance(program, str):
|
|
self.source_name = "<str>"
|
|
self.start_line = 1
|
|
self.start_column = 0
|
|
self.source = program
|
|
self.full_source = program
|
|
return
|
|
|
|
self.source_name = inspect.getsourcefile(program) # type: ignore
|
|
lines, self.start_line = getsourcelines(program) # type: ignore
|
|
if lines:
|
|
self.start_column = len(lines[0]) - len(lines[0].lstrip())
|
|
else:
|
|
self.start_column = 0
|
|
if self.start_column and lines:
|
|
self.source = "\n".join([l[self.start_column :].rstrip() for l in lines])
|
|
else:
|
|
self.source = "".join(lines)
|
|
try:
|
|
# It will cause a problem when running in Jupyter Notebook.
|
|
# `mod` will be <module '__main__'>, which is a built-in module
|
|
# and `getsource` will throw a TypeError
|
|
mod = inspect.getmodule(program)
|
|
if mod:
|
|
self.full_source = inspect.getsource(mod)
|
|
else:
|
|
self.full_source = self.source
|
|
except TypeError:
|
|
# It's a work around for Jupyter problem.
|
|
# Since `findsource` is an internal API of inspect, we just use it
|
|
# as a fallback method.
|
|
src, _ = inspect.findsource(program) # type: ignore
|
|
self.full_source = "".join(src)
|
|
|
|
def as_ast(self) -> doc.AST:
|
|
"""Parse the source code into AST.
|
|
|
|
Returns
|
|
-------
|
|
res : doc.AST
|
|
The AST of source code.
|
|
"""
|
|
return doc.parse(self.source)
|
|
|
|
|
|
_getfile = inspect.getfile # pylint: disable=invalid-name
|
|
_findsource = inspect.findsource # pylint: disable=invalid-name
|
|
|
|
|
|
def _patched_inspect_getfile(obj):
|
|
"""Work out which source or compiled file an object was defined in."""
|
|
if not inspect.isclass(obj):
|
|
return _getfile(obj)
|
|
mod = getattr(obj, "__module__", None)
|
|
if mod is not None:
|
|
file = getattr(sys.modules[mod], "__file__", None)
|
|
if file is not None:
|
|
return file
|
|
for _, member in inspect.getmembers(obj):
|
|
if inspect.isfunction(member):
|
|
if obj.__qualname__ + "." + member.__name__ == member.__qualname__:
|
|
return inspect.getfile(member)
|
|
raise TypeError(f"Source for {obj!r} not found")
|
|
|
|
|
|
def findsource(obj):
|
|
"""Return the entire source file and starting line number for an object."""
|
|
import linecache # pylint: disable=import-outside-toplevel
|
|
|
|
if not inspect.isclass(obj):
|
|
return _findsource(obj)
|
|
|
|
file = inspect.getsourcefile(obj)
|
|
if file:
|
|
linecache.checkcache(file)
|
|
else:
|
|
file = inspect.getfile(obj)
|
|
if not (file.startswith("<") and file.endswith(">")):
|
|
raise OSError("source code not available")
|
|
|
|
module = inspect.getmodule(obj, file)
|
|
if module:
|
|
lines = linecache.getlines(file, module.__dict__)
|
|
else:
|
|
lines = linecache.getlines(file)
|
|
if not lines:
|
|
raise OSError("could not get source code")
|
|
qual_names = obj.__qualname__.replace(".<locals>", "<locals>").split(".")
|
|
in_comment = 0
|
|
scope_stack = []
|
|
indent_info = {}
|
|
for i, line in enumerate(lines):
|
|
n_comment = line.count('"""')
|
|
if n_comment:
|
|
# update multi-line comments status
|
|
in_comment = in_comment ^ (n_comment & 1)
|
|
continue
|
|
if in_comment:
|
|
# skip lines within multi-line comments
|
|
continue
|
|
indent = len(line) - len(line.lstrip())
|
|
tokens = line.split()
|
|
if len(tokens) > 1:
|
|
name = None
|
|
if tokens[0] == "def":
|
|
name = tokens[1].split(":")[0].split("(")[0] + "<locals>"
|
|
elif tokens[0] == "class":
|
|
name = tokens[1].split(":")[0].split("(")[0]
|
|
# pop scope if we are less indented
|
|
while scope_stack and indent_info[scope_stack[-1]] >= indent:
|
|
scope_stack.pop()
|
|
if name:
|
|
scope_stack.append(name)
|
|
indent_info[name] = indent
|
|
if scope_stack == qual_names:
|
|
return lines, i
|
|
|
|
raise OSError("could not find class definition")
|
|
|
|
|
|
def getsourcelines(obj):
|
|
"""Extract the block of code at the top of the given list of lines."""
|
|
obj = inspect.unwrap(obj)
|
|
lines, l_num = findsource(obj)
|
|
return inspect.getblock(lines[l_num:]), l_num + 1
|
|
|
|
|
|
inspect.getfile = _patched_inspect_getfile
|
|
|
|
|
|
def _format_source_snippet(
|
|
source_lines: list,
|
|
lineno: int,
|
|
col_offset: int,
|
|
end_lineno: int,
|
|
end_col_offset: int,
|
|
) -> str:
|
|
"""Format a source code snippet with a column/span marker.
|
|
|
|
Renders every source line spanned by the diagnostic (``lineno`` through
|
|
``end_lineno``, inclusive) with a per-line gutter, followed by a caret
|
|
underline covering the offending span. For a single-line span
|
|
(``end_lineno == lineno``) only the columns ``col_offset`` (inclusive)
|
|
through ``end_col_offset`` (exclusive) are underlined; for a multi-line
|
|
span the underline covers the start column to end-of-line on the first
|
|
line, the full text of interior lines, and the start of the final line up
|
|
to ``end_col_offset``.
|
|
|
|
Parameters
|
|
----------
|
|
source_lines : list of str
|
|
Lines of the source code.
|
|
|
|
lineno : int
|
|
1-based starting line number in the source.
|
|
|
|
col_offset : int
|
|
1-based starting column (inclusive) on ``lineno``.
|
|
|
|
end_lineno : int
|
|
1-based ending line number in the source (>= ``lineno``).
|
|
|
|
end_col_offset : int
|
|
1-based ending column (exclusive) on ``end_lineno``.
|
|
|
|
Returns
|
|
-------
|
|
snippet : str
|
|
Formatted source snippet with caret-marker line(s).
|
|
"""
|
|
if end_lineno < lineno:
|
|
end_lineno = lineno
|
|
|
|
# Determine the gutter width so that all line numbers line up.
|
|
header_width = len(f" {end_lineno} ")
|
|
no_line_header = " " * header_width
|
|
|
|
parts = [f"{no_line_header}| "]
|
|
for cur_lineno in range(lineno, end_lineno + 1):
|
|
idx = cur_lineno - 1
|
|
if not 0 <= idx < len(source_lines):
|
|
continue
|
|
line_text = source_lines[idx].rstrip("\n")
|
|
line_header = f" {cur_lineno} ".rjust(header_width)
|
|
|
|
# Compute the underline span [start_col, stop_col) for this line.
|
|
start_col = col_offset if cur_lineno == lineno else 1
|
|
stop_col = end_col_offset if cur_lineno == end_lineno else len(line_text) + 1
|
|
|
|
marker = ""
|
|
for i in range(1, len(line_text) + 1):
|
|
if start_col <= i < stop_col:
|
|
marker += "^"
|
|
else:
|
|
marker += " "
|
|
parts.append(f"{line_header}| {line_text}")
|
|
parts.append(f"{no_line_header}| {marker}")
|
|
|
|
if len(parts) == 1:
|
|
return ""
|
|
return "\n".join(parts)
|
|
|
|
|
|
class Diagnostics:
|
|
"""Diagnostics class for error reporting in parser.
|
|
|
|
Formats parse errors with source location context and raises directly,
|
|
without going through DiagnosticContext.
|
|
|
|
Parameters
|
|
----------
|
|
source : Source
|
|
The source code.
|
|
"""
|
|
|
|
source: Source
|
|
|
|
def __init__(self, source: Source):
|
|
self.source = source
|
|
|
|
def error(self, node: doc.AST, message: str) -> None:
|
|
"""Emit a diagnostic error by raising with source location context.
|
|
|
|
Parameters
|
|
----------
|
|
node : doc.AST
|
|
The node with diagnostic error.
|
|
|
|
message : str
|
|
The diagnostic message.
|
|
"""
|
|
lineno = getattr(node, "lineno", 1)
|
|
col_offset = getattr(node, "col_offset", self.source.start_column)
|
|
end_lineno = getattr(node, "end_lineno", lineno)
|
|
end_col_offset = getattr(node, "end_col_offset", col_offset)
|
|
lineno += self.source.start_line - 1
|
|
end_lineno += self.source.start_line - 1
|
|
col_offset += self.source.start_column + 1
|
|
end_col_offset += self.source.start_column + 1
|
|
|
|
source_lines = self.source.full_source.splitlines(keepends=True)
|
|
snippet = _format_source_snippet(
|
|
source_lines, lineno, col_offset, end_lineno, end_col_offset
|
|
)
|
|
|
|
location = f"{self.source.source_name}:{lineno}:{col_offset}"
|
|
formatted = f"error: {message}\n --> {location}\n{snippet}"
|
|
raise DiagnosticError(formatted)
|