diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py index 6e769ce90..5623f3d11 100644 --- a/dvc/dvcfile.py +++ b/dvc/dvcfile.py @@ -2,11 +2,11 @@ import collections import contextlib import logging import os -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, Any, List, Union from voluptuous import MultipleInvalid -from dvc.exceptions import DvcException +from dvc.exceptions import DvcException, PrettyDvcException from dvc.parsing.versions import LOCKFILE_VERSION, SCHEMA_KWD from dvc.stage import serialize from dvc.stage.exceptions import ( @@ -19,6 +19,7 @@ from dvc.types import AnyPath from dvc.utils import relpath from dvc.utils.collections import apply_diff from dvc.utils.serialize import ( + YAMLFileCorruptedError, dump_yaml, modify_yaml, parse_yaml, @@ -26,6 +27,8 @@ from dvc.utils.serialize import ( ) if TYPE_CHECKING: + from ruamel.yaml import StreamMark + from dvc.repo import Repo logger = logging.getLogger(__name__) @@ -95,6 +98,78 @@ def check_dvcfile_path(repo, path): raise FileIsGitIgnored(relpath(path), True) +class YAMLSyntaxError(PrettyDvcException, YAMLFileCorruptedError): + def __init__(self, path, yaml_text, exc): + self.path = path + self.yaml_text = yaml_text + self.exc = exc + super().__init__(self.path) + + def __pretty_exc__(self, **kwargs: Any): + from rich.syntax import Syntax + from rich.text import Text + from ruamel.yaml.error import MarkedYAMLError + + from dvc.ui import ui + + exc = self.exc.__cause__ + + if not isinstance(exc, MarkedYAMLError): + raise ValueError("nothing to pretty-print here. :)") + + source = self.yaml_text.splitlines() + + def prepare_linecol(mark: "StreamMark") -> str: + return f"in line {mark.line + 1}, column {mark.column + 1}" + + def prepare_message(message: str, mark: "StreamMark" = None) -> "Text": + return Text.assemble( + message.capitalize(), + f", {prepare_linecol(mark)}" if mark else "", + style="bold", + ) + + def prepare_code(mark: "StreamMark") -> "Syntax": + line = mark.line + 1 + code = "" if line > len(source) else source[line - 1] + return Syntax( + code, + "yaml", + start_line=line, + theme="ansi_dark", + word_wrap=True, + line_numbers=True, + indent_guides=True, + ) + + lines: List[object] = [] + if hasattr(exc, "context"): + if exc.context_mark is not None: + lines.append( + prepare_message(str(exc.context), exc.context_mark) + ) + if exc.context_mark is not None and ( + exc.problem is None + or exc.problem_mark is None + or exc.context_mark.name != exc.problem_mark.name + or exc.context_mark.line != exc.problem_mark.line + or exc.context_mark.column != exc.problem_mark.column + ): + lines.extend([prepare_code(exc.context_mark), ""]) + if exc.problem is not None: + lines.append( + prepare_message(str(exc.problem), exc.problem_mark) + ) + if exc.problem_mark is not None: + lines.append(prepare_code(exc.problem_mark)) + + if lines: + lines.insert(0, "") + lines.insert(0, f"[red]'{relpath(self.path)}' structure is corrupted.") + for message in lines: + ui.error_write(message, styled=True) + + class FileMixin: SCHEMA = None @@ -144,6 +219,8 @@ class FileMixin: # 2. filename is not a DVC file # 3. path doesn't represent a regular file # 4. when the file is git ignored + from dvc.utils.serialize import ParseError + if not self.exists(): dvc_ignored = self.repo.dvcignore.is_ignored_file(self.path) raise StageFileDoesNotExistError( @@ -157,7 +234,13 @@ class FileMixin: self._check_gitignored() with self.repo.fs.open(self.path, encoding="utf-8") as fd: stage_text = fd.read() - d = parse_yaml(stage_text, self.path) + + try: + d = parse_yaml(stage_text, self.path) + except ParseError as exc: + cause = exc.__cause__ + raise YAMLSyntaxError(self.path, stage_text, exc) from cause + return self.validate(d, self.relpath), stage_text @classmethod diff --git a/dvc/exceptions.py b/dvc/exceptions.py index 6d5db86f4..3c2b84af4 100644 --- a/dvc/exceptions.py +++ b/dvc/exceptions.py @@ -364,3 +364,8 @@ class CircularImportError(DvcException): f"'{dep}' contains invalid circular import. " f"DVC repo '{a}' already imports from '{b}'." ) + + +class PrettyDvcException(DvcException): + def __pretty_exc__(self, **kwargs): + """Print prettier exception message.""" diff --git a/dvc/logger.py b/dvc/logger.py index cd32069ef..849e9b360 100644 --- a/dvc/logger.py +++ b/dvc/logger.py @@ -142,9 +142,22 @@ class LoggerHandler(logging.StreamHandler): super().handleError(record) raise LoggingException(record) + def emit_pretty_exception(self, exc, verbose: bool = False): + return exc.__pretty_exc__(verbose=verbose) + def emit(self, record): """Write to Tqdm's stream so as to not break progress-bars""" try: + if record.exc_info: + _, exc, *_ = record.exc_info + if hasattr(exc, "__pretty_exc__"): + try: + self.emit_pretty_exception(exc, verbose=_is_verbose()) + if not _is_verbose(): + return + except Exception: # noqa, pylint: disable=broad-except + pass + msg = self.format(record) Tqdm.write( msg, file=self.stream, end=getattr(self, "terminator", "\n")