diff --git a/guppylang/decorator.py b/guppylang/decorator.py index 645c5f06..90d79e9b 100644 --- a/guppylang/decorator.py +++ b/guppylang/decorator.py @@ -38,6 +38,7 @@ find_guppy_module_in_py_module, get_calling_frame, ) +from guppylang.span import SourceMap from guppylang.tys.subst import Inst from guppylang.tys.ty import NumericType @@ -77,8 +78,12 @@ class _Guppy: # The currently-alive GuppyModules, associated with a Python file/module _modules: dict[ModuleIdentifier, GuppyModule] + # Storage for source code that has been read by the compiler + _sources: SourceMap + def __init__(self) -> None: self._modules = {} + self._sources = SourceMap() @overload def __call__(self, arg: PyFunc) -> RawFunctionDef: ... @@ -300,7 +305,7 @@ def custom( mod = module or self.get_module() def dec(f: PyFunc) -> RawCustomFunctionDef: - func_ast, docstring = parse_py_func(f) + func_ast, docstring = parse_py_func(f, self._sources) if not has_empty_body(func_ast): raise GuppyError( "Body of custom function declaration must be empty", @@ -360,7 +365,9 @@ def constant( ) -> RawConstDef: """Adds a constant to a module, backed by a `hugr.val.Value`.""" module = module or self.get_module() - type_ast = _parse_expr_string(ty, f"Not a valid Guppy type: `{ty}`") + type_ast = _parse_expr_string( + ty, f"Not a valid Guppy type: `{ty}`", self._sources + ) defn = RawConstDef(DefId.fresh(module), name, None, type_ast, value) module.register_def(defn) return defn @@ -375,7 +382,9 @@ def extern( ) -> RawExternDef: """Adds an extern symbol to a module.""" module = module or self.get_module() - type_ast = _parse_expr_string(ty, f"Not a valid Guppy type: `{ty}`") + type_ast = _parse_expr_string( + ty, f"Not a valid Guppy type: `{ty}`", self._sources + ) defn = RawExternDef( DefId.fresh(module), name, None, symbol or name, constant, type_ast ) @@ -444,7 +453,7 @@ def registered_modules(self) -> KeysView[ModuleIdentifier]: guppy = _Guppy() -def _parse_expr_string(ty_str: str, parse_err: str) -> ast.expr: +def _parse_expr_string(ty_str: str, parse_err: str, sources: SourceMap) -> ast.expr: """Helper function to parse expressions that are provided as strings. Tries to infer the source location were the given string was defined by inspecting @@ -460,6 +469,7 @@ def _parse_expr_string(ty_str: str, parse_err: str) -> ast.expr: if caller_frame := get_calling_frame(): info = inspect.getframeinfo(caller_frame) if caller_module := inspect.getmodule(caller_frame): + sources.add_file(info.filename) source_lines, _ = inspect.getsourcelines(caller_module) source = "".join(source_lines) annotate_location(expr_ast, source, info.filename, 0) diff --git a/guppylang/definition/common.py b/guppylang/definition/common.py index 73e8278b..0a11d882 100644 --- a/guppylang/definition/common.py +++ b/guppylang/definition/common.py @@ -8,6 +8,8 @@ from hugr.build.dfg import DefinitionBuilder, OpVar from hugr.ext import Package +from guppylang.span import SourceMap + if TYPE_CHECKING: from guppylang.checker.core import Globals from guppylang.compiler.core import CompiledGlobals @@ -92,7 +94,7 @@ class ParsableDef(Definition): """ @abstractmethod - def parse(self, globals: "Globals") -> ParsedDef: + def parse(self, globals: "Globals", sources: SourceMap) -> ParsedDef: """Performs parsing and validation, returning a definition that can be checked. The provided globals contain all other raw definitions that have been defined. diff --git a/guppylang/definition/const.py b/guppylang/definition/const.py index d7479425..1f691317 100644 --- a/guppylang/definition/const.py +++ b/guppylang/definition/const.py @@ -10,6 +10,7 @@ from guppylang.compiler.core import CompiledGlobals, DFContainer from guppylang.definition.common import CompilableDef, ParsableDef from guppylang.definition.value import CompiledValueDef, ValueDef +from guppylang.span import SourceMap from guppylang.tys.parsing import type_from_ast @@ -22,7 +23,7 @@ class RawConstDef(ParsableDef): description: str = field(default="constant", init=False) - def parse(self, globals: Globals) -> "ConstDef": + def parse(self, globals: Globals, sources: SourceMap) -> "ConstDef": """Parses and checks the user-provided signature of the function.""" return ConstDef( self.id, diff --git a/guppylang/definition/custom.py b/guppylang/definition/custom.py index a3eeef5a..ee250d7b 100644 --- a/guppylang/definition/custom.py +++ b/guppylang/definition/custom.py @@ -16,6 +16,7 @@ from guppylang.definition.value import CallReturnWires, CompiledCallableDef from guppylang.error import GuppyError, InternalGuppyError from guppylang.nodes import GlobalCall +from guppylang.span import SourceMap from guppylang.tys.subst import Inst, Subst from guppylang.tys.ty import ( FuncInput, @@ -56,7 +57,7 @@ class RawCustomFunctionDef(ParsableDef): description: str = field(default="function", init=False) - def parse(self, globals: "Globals") -> "CustomFunctionDef": + def parse(self, globals: "Globals", sources: SourceMap) -> "CustomFunctionDef": """Parses and checks the user-provided signature of the custom function. The signature is optional if custom type checking logic is provided by the user. diff --git a/guppylang/definition/declaration.py b/guppylang/definition/declaration.py index 2352ca7b..7a77ac10 100644 --- a/guppylang/definition/declaration.py +++ b/guppylang/definition/declaration.py @@ -16,6 +16,7 @@ from guppylang.definition.value import CallableDef, CallReturnWires, CompiledCallableDef from guppylang.error import GuppyError from guppylang.nodes import GlobalCall +from guppylang.span import SourceMap from guppylang.tys.subst import Inst, Subst from guppylang.tys.ty import Type, type_to_row @@ -32,9 +33,9 @@ class RawFunctionDecl(ParsableDef): python_scope: PyScope description: str = field(default="function", init=False) - def parse(self, globals: Globals) -> "CheckedFunctionDecl": + def parse(self, globals: Globals, sources: SourceMap) -> "CheckedFunctionDecl": """Parses and checks the user-provided signature of the function.""" - func_ast, docstring = parse_py_func(self.python_func) + func_ast, docstring = parse_py_func(self.python_func, sources) ty = check_signature(func_ast, globals.with_python_scope(self.python_scope)) if not has_empty_body(func_ast): raise GuppyError( diff --git a/guppylang/definition/extern.py b/guppylang/definition/extern.py index 3950750d..a664c91a 100644 --- a/guppylang/definition/extern.py +++ b/guppylang/definition/extern.py @@ -9,6 +9,7 @@ from guppylang.compiler.core import CompiledGlobals, DFContainer from guppylang.definition.common import CompilableDef, ParsableDef from guppylang.definition.value import CompiledValueDef, ValueDef +from guppylang.span import SourceMap from guppylang.tys.parsing import type_from_ast @@ -22,7 +23,7 @@ class RawExternDef(ParsableDef): description: str = field(default="extern", init=False) - def parse(self, globals: Globals) -> "ExternDef": + def parse(self, globals: Globals, sources: SourceMap) -> "ExternDef": """Parses and checks the user-provided signature of the function.""" return ExternDef( self.id, diff --git a/guppylang/definition/function.py b/guppylang/definition/function.py index 3920fe32..8187566f 100644 --- a/guppylang/definition/function.py +++ b/guppylang/definition/function.py @@ -26,6 +26,7 @@ from guppylang.error import GuppyError from guppylang.ipython_inspect import find_ipython_def, is_running_ipython from guppylang.nodes import GlobalCall +from guppylang.span import SourceMap from guppylang.tys.subst import Inst, Subst from guppylang.tys.ty import FunctionType, Type, type_to_row @@ -53,9 +54,9 @@ class RawFunctionDef(ParsableDef): description: str = field(default="function", init=False) - def parse(self, globals: Globals) -> "ParsedFunctionDef": + def parse(self, globals: Globals, sources: SourceMap) -> "ParsedFunctionDef": """Parses and checks the user-provided signature of the function.""" - func_ast, docstring = parse_py_func(self.python_func) + func_ast, docstring = parse_py_func(self.python_func, sources) ty = check_signature(func_ast, globals.with_python_scope(self.python_scope)) if ty.parametrized: raise GuppyError( @@ -220,7 +221,7 @@ def compile_inner(self, globals: CompiledGlobals) -> None: compile_global_func_def(self, self.func_def, globals) -def parse_py_func(f: PyFunc) -> tuple[ast.FunctionDef, str | None]: +def parse_py_func(f: PyFunc, sources: SourceMap) -> tuple[ast.FunctionDef, str | None]: source_lines, line_offset = inspect.getsourcelines(f) source = "".join(source_lines) # Lines already have trailing \n's source = textwrap.dedent(source) @@ -234,10 +235,12 @@ def parse_py_func(f: PyFunc) -> tuple[ast.FunctionDef, str | None]: defn = find_ipython_def(func_ast.name) if defn is not None: file = f"<{defn.cell_name}>" + sources.add_file(file, source) else: file = inspect.getsourcefile(f) - if file is None: - raise GuppyError("Couldn't determine source file for function") + if file is None: + raise GuppyError("Couldn't determine source file for function") + sources.add_file(file) annotate_location(func_ast, source, file, line_offset) if not isinstance(func_ast, ast.FunctionDef): raise GuppyError("Expected a function definition", func_ast) diff --git a/guppylang/definition/struct.py b/guppylang/definition/struct.py index 521b5549..f1573322 100644 --- a/guppylang/definition/struct.py +++ b/guppylang/definition/struct.py @@ -25,6 +25,7 @@ from guppylang.definition.ty import TypeDef from guppylang.error import GuppyError, InternalGuppyError from guppylang.ipython_inspect import find_ipython_def, is_running_ipython +from guppylang.span import SourceMap from guppylang.tys.arg import Argument from guppylang.tys.param import Parameter, check_all_args from guppylang.tys.parsing import type_from_ast @@ -63,9 +64,9 @@ def __getitem__(self, item: Any) -> "RawStructDef": """ return self - def parse(self, globals: Globals) -> "ParsedStructDef": + def parse(self, globals: Globals, sources: SourceMap) -> "ParsedStructDef": """Parses the raw class object into an AST and checks that it is well-formed.""" - cls_def = parse_py_class(self.python_class) + cls_def = parse_py_class(self.python_class, sources) if cls_def.keywords: raise GuppyError("Unexpected keyword", cls_def.keywords[0]) @@ -232,7 +233,7 @@ def compile(self, args: list[Wire]) -> list[Wire]: return [constructor_def] -def parse_py_class(cls: type) -> ast.ClassDef: +def parse_py_class(cls: type, sources: SourceMap) -> ast.ClassDef: """Parses a Python class object into an AST.""" # If we are running IPython, `inspect.getsourcelines` works only for builtins # (guppy stdlib), but not for most/user-defined classes - see: @@ -254,6 +255,8 @@ def parse_py_class(cls: type) -> ast.ClassDef: file = inspect.getsourcefile(cls) if file is None: raise GuppyError("Couldn't determine source file for class") + # Store the source file in our cache + sources.add_file(file) annotate_location(cls_ast, source, file, line_offset) if not isinstance(cls_ast, ast.ClassDef): raise GuppyError("Expected a class definition", cls_ast) diff --git a/guppylang/module.py b/guppylang/module.py index c4a2ced4..4befecd2 100644 --- a/guppylang/module.py +++ b/guppylang/module.py @@ -29,6 +29,7 @@ from guppylang.definition.ty import TypeDef from guppylang.error import GuppyError, pretty_errors from guppylang.experimental import enable_experimental_features +from guppylang.span import SourceMap PyClass = type PyFunc = Callable[..., Any] @@ -73,6 +74,9 @@ class GuppyModule: # `_register_buffered_instance_funcs` is called. This way, we can associate _instance_func_buffer: dict[str, RawDef] | None + # Storage for source code that has been read by the compiler + _sources: SourceMap + def __init__(self, name: str, import_builtins: bool = True): self.name = name self._globals = Globals({}, {}, {}, {}) @@ -86,6 +90,10 @@ def __init__(self, name: str, import_builtins: bool = True): self._raw_type_defs = {} self._checked_defs = {} + from guppylang.decorator import guppy + + self._sources = guppy._sources + # Import builtin module if import_builtins: import guppylang.prelude.builtins as builtins @@ -249,14 +257,15 @@ def checked(self) -> bool: def compiled(self) -> bool: return self._compiled - @staticmethod def _check_defs( - raw_defs: Mapping[DefId, RawDef], globals: Globals + self, raw_defs: Mapping[DefId, RawDef], globals: Globals ) -> dict[DefId, CheckedDef]: """Helper method to parse and check raw definitions.""" raw_globals = globals | Globals(dict(raw_defs), {}, {}, {}) parsed = { - def_id: defn.parse(raw_globals) if isinstance(defn, ParsableDef) else defn + def_id: defn.parse(raw_globals, self._sources) + if isinstance(defn, ParsableDef) + else defn for def_id, defn in raw_defs.items() } parsed_globals = globals | Globals(dict(parsed), {}, {}, {}) diff --git a/guppylang/span.py b/guppylang/span.py index e92d9688..743e289a 100644 --- a/guppylang/span.py +++ b/guppylang/span.py @@ -1,6 +1,7 @@ """Source spans representing locations in the code being compiled.""" import ast +import linecache from dataclasses import dataclass from typing import TypeAlias @@ -80,3 +81,31 @@ def to_span(x: ToSpan) -> Span: x.end_col_offset or x.col_offset, ) return Span(start, end) + + +#: List of source lines in a file +SourceLines: TypeAlias = list[str] + + +class SourceMap: + """Map holding the source code for all files accessed by the compiler. + + Can be used to look up the source code associated with a span. + """ + + sources: dict[str, SourceLines] + + def __init__(self) -> None: + self.sources = {} + + def add_file(self, file: str, content: str | None = None) -> None: + """Registers a new source file.""" + if content is None: + self.sources[file] = [line.rstrip() for line in linecache.getlines(file)] + else: + self.sources[file] = content.splitlines(keepends=False) + + def span_lines(self, span: Span, prefix_lines: int = 0) -> list[str]: + return self.sources[span.file][ + span.start.line - prefix_lines - 1 : span.end.line + ]