fix: improve error output when parsing fails

src/puya/compile.py

@@ -17,15 +17,21 @@
 from puya.codegen.builder import compile_ir_to_teal
 from puya.codegen.emitprogram import CompiledContract
 from puya.context import CompileContext
-from puya.errors import Errors, InternalError, ParseError, log_exceptions
+from puya.errors import ErrorExitCode, Errors, InternalError, ParseError, log_exceptions
 from puya.ir.context import IRBuildContext
 from puya.ir.destructure.main import destructure_ssa
 from puya.ir.main import build_embedded_ir, build_ir
 from puya.ir.optimize.dead_code_elimination import remove_unused_subroutines
 from puya.ir.optimize.main import optimize_contract_ir
 from puya.ir.to_text_visitor import output_contract_ir_to_path
 from puya.options import PuyaOptions
-from puya.parse import EMBEDDED_MODULES, TYPESHED_PATH, ParseSource, parse_and_typecheck
+from puya.parse import (
+    EMBEDDED_MODULES,
+    TYPESHED_PATH,
+    ParseSource,
+    SourceLocation,
+    parse_and_typecheck,
+)
 from puya.utils import attrs_extend, determine_out_dir, make_path_relative_to_cwd
 
 logger = structlog.get_logger(__name__)
@@ -54,11 +60,10 @@ def parse_with_mypy(puya_options: PuyaOptions) -> CompileContext:
     if read_source is None:
         raise InternalError("parse_results.manager.errors.read_source is None")
 
-    errors = Errors()
     context = CompileContext(
         options=puya_options,
         parse_result=parse_result,
-        errors=errors,
+        errors=Errors(read_source),
         read_source=read_source,
     )
 
@@ -187,10 +192,14 @@ def log_options(puya_options: PuyaOptions) -> None:
 def compile_to_teal(puya_options: PuyaOptions) -> None:
     """Drive the actual core compilation step."""
     log_options(puya_options)
-    context = parse_with_mypy(puya_options)
-    awst = transform_ast(context)
+    try:
+        context = parse_with_mypy(puya_options)
+    except ParseError as ex:
+        _log_parse_errors(ex)
+        sys.exit(ErrorExitCode.code)
 
-    with log_exceptions(context.errors):
+    with log_exceptions(context.errors, exit_on_failure=True):
+        awst = transform_ast(context)
         compiled_contracts_by_source_path = awst_to_teal(context, awst)
         write_teal_to_output(context, compiled_contracts_by_source_path)
 
@@ -200,7 +209,7 @@ def write_teal_to_output(
 ) -> None:
     if contracts is None:
         logger.error("Build failed")
-        sys.exit(1)
+        sys.exit(ErrorExitCode.code)
     elif not contracts:
         logger.error("No contracts discovered in any source files")
     else:
@@ -267,3 +276,20 @@ def write_contract_files(base_path: Path, compiled_contract: CompiledContract) -
         output_text = "\n".join(src)
         logger.info(f"Writing {make_path_relative_to_cwd(output_path)}")
         output_path.write_text(output_text, encoding="utf-8")
+
+
+def _log_parse_errors(ex: ParseError) -> None:
+    location: SourceLocation | None = None
+    related_errors = list[str]()
+    for error in ex.errors:
+        if error.location:  # align related error messages
+            if related_errors:
+                logger.error(
+                    related_errors[0], location=location, related_lines=related_errors[1:]
+                )
+            related_errors = [error.message]
+            location = error.location
+        else:
+            related_errors.append(error.message)
+    if related_errors:
+        logger.error(related_errors[0], location=location, related_lines=related_errors[1:])

src/puya/errors.py

@@ -1,29 +1,71 @@
 import contextlib
+import enum
 import logging
 import sys
 import traceback
 import typing
 import typing as t
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
+from pathlib import Path
 
+import attrs
 import structlog
 
 from puya.parse import SourceLocation
 
 logger = structlog.get_logger()
+_VALID_SEVERITY = ("error", "note", "warning")
+
+
+class ErrorExitCode(enum.IntEnum):
+    code = 1
+    internal = 2
+    unexpected = 3
 
 
 class PuyaError(Exception):
     def __init__(self, msg: str, location: SourceLocation | None = None):
-        self._msg = msg
         super().__init__(msg)
+        self._msg = msg
         self.location = location
 
 
+@attrs.frozen(kw_only=True)
+class MyPyErrorData:
+    message: str
+    severity: str | None = None
+    location: SourceLocation | None = None
+
+    @classmethod
+    def parse(cls, error_str: str) -> "MyPyErrorData":
+        error_parts = error_str.split(":", maxsplit=3)
+        if len(error_parts) == 4:  # parse error that might contain file and line details
+            file_str, line_str, severity, message = error_parts
+            severity = severity.strip()
+            if Path(file_str).exists() and severity in _VALID_SEVERITY:
+                try:
+                    line = int(line_str)
+                except ValueError:
+                    pass
+                else:
+                    location = SourceLocation(file=file_str, line=line)
+                    return cls(message=message[1:], severity=severity, location=location)
+            error_parts = error_str.split(":", maxsplit=1)
+        if len(error_parts) == 2:  # otherwise attempt to parse severity and message
+            severity, message = error_parts
+            severity = severity.strip()
+            if severity in _VALID_SEVERITY:
+                return cls(message=message[1:], severity=severity)
+        # fallback to just the error message
+        return cls(message=error_str)
+
+
 class ParseError(Exception):
+    """Encapsulate parse/type errors from MyPy"""
+
     def __init__(self, errors: list[str]):
-        self.errors = errors
         super().__init__("\n".join(errors))
+        self.errors = list(map(MyPyErrorData.parse, errors))
 
 
 class InternalError(PuyaError):
@@ -41,15 +83,41 @@ def __init__(self, location: SourceLocation | None, msg: str | None = None):
         super().__init__(msg or "TODO: support this thing", location=location)
 
 
+def _get_pretty_source(file_source: list[str], location: SourceLocation) -> list[str]:
+    source_line = file_source[location.line - 1]
+    column = location.column
+    end_column = len(source_line)
+    if (location.end_line is None or location.end_line == location.line) and location.end_column:
+        end_column = location.end_column
+    source_line_expanded = source_line.expandtabs()
+
+    # Shifts column after tab expansion
+    column = len(source_line[:column].expandtabs())
+    end_column = len(source_line[:end_column].expandtabs())
+
+    marker = "^"
+    if end_column > column:
+        marker += "~" * (end_column - column - 1)
+    return [
+        source_line_expanded,
+        " " * column + marker,
+    ]
+
+
 class Errors:
-    def __init__(self) -> None:
+    def __init__(self, read_source: Callable[[str], list[str] | None]) -> None:
         self.num_errors = 0
         self.num_warnings = 0
+        self.read_source = read_source
 
     def _report(self, log_level: int, msg: str, location: SourceLocation | None) -> None:
         kwargs: dict[str, typing.Any] = {}
         if location:
             kwargs["location"] = location
+            if log_level == logging.ERROR:
+                file_source = self.read_source(location.file)
+                if file_source and location.line <= len(file_source):
+                    kwargs["related_lines"] = _get_pretty_source(file_source, location)
         logger.log(log_level, msg, **kwargs)
 
     def error(self, msg: str, location: SourceLocation | None) -> None:
@@ -64,7 +132,7 @@ def note(self, msg: str, location: SourceLocation | None) -> None:
         self._report(logging.INFO, msg, location)
 
 
-def crash_report(location: SourceLocation | None) -> t.Never:
+def crash_report(location: SourceLocation | None, exit_code: ErrorExitCode) -> t.Never:
     # Adapted from report_internal_error in mypy
     err = sys.exc_info()[1]
     tb = traceback.extract_stack()[:-4]
@@ -80,20 +148,25 @@ def crash_report(location: SourceLocation | None) -> t.Never:
         print(s.rstrip("\n"))
     if location:
         print(f"{location.file}:{location.line}: {type(err).__name__}: {err}")
-    raise SystemExit(2)
+    raise SystemExit(exit_code.value)
 
 
 @contextlib.contextmanager
 def log_exceptions(
-    errors: Errors, fallback_location: SourceLocation | None = None
+    errors: Errors,
+    fallback_location: SourceLocation | None = None,
+    *,
+    exit_on_failure: bool = False,
 ) -> Iterator[None]:
     try:
         yield
     except CodeError as ex:
         errors.error(str(ex), location=ex.location or fallback_location)
+        if exit_on_failure:
+            raise SystemExit(ErrorExitCode.code) from ex
     except InternalError as ex:
         errors.error(f"FATAL {ex!s}", location=ex.location or fallback_location)
-        crash_report(ex.location or fallback_location)
+        crash_report(ex.location or fallback_location, ErrorExitCode.internal)
     except Exception as ex:
         errors.error(f"UNEXPECTED {ex!s}", location=fallback_location)
-        crash_report(fallback_location)
+        crash_report(fallback_location, ErrorExitCode.unexpected)

src/puya/ir/context.py

@@ -5,7 +5,7 @@
 
 import puya.awst.nodes as awst_nodes
 from puya.context import CompileContext
-from puya.errors import CodeError, InternalError, crash_report
+from puya.errors import CodeError, InternalError, log_exceptions
 from puya.ir.models import Subroutine
 from puya.parse import SourceLocation
 from puya.utils import attrs_extend
@@ -114,16 +114,8 @@ class IRBuildContextWithFallback(IRBuildContext):
     @contextlib.contextmanager
     def log_exceptions(self, fallback_location: SourceLocation | None = None) -> Iterator[None]:
         fallback_location = fallback_location or self.default_fallback
-        try:
+        with log_exceptions(self.errors, fallback_location):
             yield
-        except CodeError as ex:
-            self.errors.error(str(ex), location=ex.location or fallback_location)
-        except InternalError as ex:
-            self.errors.error(f"FATAL {ex!s}", location=ex.location or fallback_location)
-            crash_report(ex.location or fallback_location)
-        except Exception as ex:
-            self.errors.error(f"UNEXPECTED {ex!s}", location=fallback_location)
-            crash_report(fallback_location)
 
 
 @attrs.frozen(kw_only=True)

src/puya/logging_config.py

@@ -82,28 +82,39 @@ def __call__(
         _name: str,
         event_dict: structlog.typing.EventDict,
     ) -> str:
-        sio = StringIO()
-
-        location = event_dict.pop("location", None)
-        if location:
-            sio.write(self._styles.logger_name)
-            sio.write(self._location_as_link(location))
-            sio.write(" ")
-            sio.write(self._styles.reset)
-
-        level = event_dict.pop("level", None)
-        if level is not None:
-            sio.write(self.level_to_color.get(level, ""))
-            sio.write(level)
-            sio.write(": ")
-            sio.write(self._styles.reset)
-
         # force event to str for compatibility with standard library
         event = event_dict.pop("event", None)
         if not isinstance(event, str):
             event = str(event)
-
-        sio.write(event)
+        lines = [event]
+        related_errors = event_dict.pop("related_lines", None)
+        if related_errors:
+            assert isinstance(related_errors, list)
+            lines.extend(related_errors)
+        location: SourceLocation | None = event_dict.pop("location", None)
+        location_as_link = self._location_as_link(location) if location else ""
+        level = event_dict.pop("level", "")
+
+        align_related_lines = " " * (len(location_as_link) + 1 + len(level) + 2)
+        sio = StringIO()
+        for idx, line in enumerate(lines):
+            if idx:
+                sio.write("\n")
+                sio.write(align_related_lines)
+            else:
+                if location:
+                    sio.write(self._styles.logger_name)
+                    location_link = self._location_as_link(location)
+                    sio.write(location_link)
+                    sio.write(" ")
+                    sio.write(self._styles.reset)
+
+                if level:
+                    sio.write(self.level_to_color.get(level, ""))
+                    sio.write(level)
+                    sio.write(": ")
+                    sio.write(self._styles.reset)
+            sio.write(line)
 
         stack = event_dict.pop("stack", None)
         exc = event_dict.pop("exception", None)

tests/test_cli.py

@@ -1,9 +1,15 @@
+import os
 import shutil
 import subprocess
 from pathlib import Path
 
 from tests import EXAMPLES_DIR, TEST_CASES_DIR, VCS_ROOT
 
+ENV_WITH_NO_COLOR = dict(os.environ) | {
+    "NO_COLOR": "1",  # disable colour output
+    "PYTHONUTF8": "1",  # force utf8 on windows
+}
+
 
 def run_puyapy(
     args: list[str | Path], *, check: bool = True, text: bool = True
@@ -23,6 +29,7 @@ def run_puyapy(
         stderr=subprocess.STDOUT,
         cwd=VCS_ROOT,
         check=False,
+        env=ENV_WITH_NO_COLOR,
     )
     if check:
         assert result.returncode == 0, result.stdout