Materialize/misc/python/materialize/output_consistency/output/base_output_printer.py

# Copyright Materialize, Inc. and contributors. All rights reserved.
#
# Use of this software is governed by the Business Source License
# included in the LICENSE file at the root of this repository.
#
# As of the Change Date specified in that file, in accordance with
# the Business Source License, use of this software will be governed
# by the Apache License, Version 2.0.

import re
from enum import Enum

from materialize.output_consistency.output.format_constants import (
    COMMENT_PREFIX,
    OUTPUT_LINE_SEPARATOR,
    SECTION_COLLAPSED_PREFIX,
    SECTION_EXPANDED_PREFIX,
)


class OutputPrinterMode(Enum):
    PRINT = 1
    COLLECT = 2


class BaseOutputPrinter:

    def __init__(self, print_mode: OutputPrinterMode = OutputPrinterMode.PRINT):
        self.print_mode = print_mode
        self.collected_output = []

    def print_empty_line(self) -> None:
        self._print_raw("")

    def start_section(self, header: str, collapsed: bool = True) -> None:
        prefix = SECTION_COLLAPSED_PREFIX if collapsed else SECTION_EXPANDED_PREFIX
        self._print_raw(f"{prefix}{header}")

    def print_separator_line(self) -> None:
        self._print_text(OUTPUT_LINE_SEPARATOR)

    def _print_executable(self, sql: str) -> None:
        self._print_raw(sql)

    def _print_text(self, text: str) -> None:
        adjusted_text = re.sub("\n", f"\n{COMMENT_PREFIX}", text)
        adjusted_text = f"{COMMENT_PREFIX}{adjusted_text}"

        self._print_raw(adjusted_text)

    def _print_raw(self, sql: str) -> None:
        if self.print_mode == OutputPrinterMode.PRINT:
            print(sql)
        elif self.print_mode == OutputPrinterMode.COLLECT:
            self.collected_output.append(sql)
        else:
            raise RuntimeError(f"Unsupported mode: {self.print_mode}")