src/sqlfmt/jinjafmt.py

import ast
import keyword
import re
from dataclasses import dataclass, field
from importlib import import_module
from itertools import chain, product
from types import ModuleType
from typing import Dict, List, MutableSet, NamedTuple, Optional, Tuple

from sqlfmt.line import Line
from sqlfmt.mode import Mode
from sqlfmt.node import Node
from sqlfmt.node_manager import NodeManager
from sqlfmt.splitter import LineSplitter


class BlackWrapper:
    """
    A thin wrapper around black. Tries to import black when
    instantiated. Provides a safe interface, format_string
    """

    PY_RESERVED_WORDS = list(keyword.kwlist)

    class StringProperties(NamedTuple):
        has_newlines: bool
        keyword_replacements: Dict[str, int]
        tilde_replacements: Dict[str, int]

    def __init__(self) -> None:
        try:
            self.black: Optional[ModuleType] = import_module("black")
        except ImportError:
            self.black = None

    def format_string(self, source_string: str, max_length: int) -> Tuple[str, bool]:
        """
        Attempt to use black to format source_string to a line_length of max_length.
        Return source_string if black isn't installed or it can't parse source_string.

        Return a tuple of the formatted string and a boolean that indicates whether
        black successfully ran on the string
        """
        BLACKENED = True
        NOT_BLACKENED = False

        if not self.black:
            return source_string, NOT_BLACKENED

        try:
            preprocessed_string, string_properties = self._preprocess_string(
                source_string
            )
        except ValueError:
            return source_string, NOT_BLACKENED

        black_mode = self.black.Mode(line_length=max_length)

        try:
            formatted_string = self.black.format_str(
                preprocessed_string, mode=black_mode
            ).rstrip()
        except ValueError:
            # the string isn't valid python.
            # Jinja allows linebreaks where python doesn't
            # so let's try again without newlines in the code
            if string_properties.has_newlines:
                try:
                    flat_code = preprocessed_string.replace("\n", " ")
                    formatted_string = self.black.format_str(
                        flat_code, mode=black_mode
                    ).rstrip()
                except ValueError:
                    # there is other jinja syntax that isn't valid python,
                    # so if this still fails, just stop trying
                    return source_string, NOT_BLACKENED
                else:
                    postprocessed_string = self._postprocess_string(
                        formatted_string, string_properties
                    )
                    return postprocessed_string, BLACKENED
            else:
                return source_string, NOT_BLACKENED
        else:
            postprocessed_string = self._postprocess_string(
                formatted_string, string_properties
            )
            return postprocessed_string, BLACKENED

    @classmethod
    def _preprocess_string(cls, source_string: str) -> Tuple[str, StringProperties]:
        """
        Takes a jinja source_string and performs some small transformations on it to
        make it valid python that black can format:
        1. Detects newlines
        2. substitutes python keywords like return for safe alternatives like return_
        3. substitutes jinja tilde operators with another binary operator

        Returns a tuple of the processed string and a NamedTuple of the stats from
        pre-processing
        """
        has_newline = True if "\n" in source_string else False

        processed_string, keyword_replacements = cls._replace_reserved_words(
            source_string=source_string
        )
        processed_string, tilde_replacements = cls._replace_tildes(
            source_string=processed_string
        )

        props = cls.StringProperties(
            has_newlines=has_newline,
            keyword_replacements=keyword_replacements,
            tilde_replacements=tilde_replacements,
        )
        return processed_string, props

    @classmethod
    def _replace_reserved_words(cls, source_string: str) -> Tuple[str, Dict[str, int]]:
        """
        Replaces python reserved words in source_string when they are used as variables
        or function names. Returns a string with the replacements made and a dict of
        the number and types of replacements made
        """
        suffixes = [r"\s*=", r"\("]
        replacements = {
            # kw patt: replacement patt, replacement
            # e.g. r"return\(": (r"return_\(", "return_(")
            f"{w}{s}": (f"{w}_{s}", f"{w}_{s[-1]}")
            for (w, s) in product(cls.PY_RESERVED_WORDS, suffixes)
        }

        # check to make sure there aren't already variables with the replacement
        # names in the source string
        preexisting_sentinels = any(
            [
                re.search(repl_patt, source_string)
                and re.search(kw_patt, source_string)
                for kw_patt, (repl_patt, _) in replacements.items()
            ]
        )
        if preexisting_sentinels:
            # abort
            raise ValueError("Cannot preprocess due to preexisting sentinels")

        # try to replace any instances of reserved words with a safe alternative
        processed_string = source_string
        keyword_replacements = {}
        for patt, (repl_patt, repl) in replacements.items():
            processed_string, n = re.subn(patt, repl, processed_string)
            if n > 0:
                keyword_replacements[repl_patt] = n

        return processed_string, keyword_replacements

    @classmethod
    def _replace_tildes(cls, source_string: str) -> Tuple[str, Dict[str, int]]:
        """
        Jinja uses ~ as the string concatenation operator, but black cannot parse the
        tilde. This method finds another operator to safely replace the tilde with,
        and returns a string with the tilde replaced and a dict with the symbol it
        was replaced with and the number of replacements made.
        """
        if "~" in source_string:
            operators = ["+", "-", "*", "/"]
            for operator in operators:
                if operator in source_string:
                    continue
                else:
                    n = source_string.count("~")
                    processed_string = source_string.replace("~", operator)
                    return processed_string, {operator: n}
            else:
                return source_string, {}
        else:
            return source_string, {}

    @classmethod
    def _postprocess_string(
        cls,
        formatted_string: str,
        string_properties: StringProperties,
    ) -> str:
        """
        Translates a formatted python string back to jinja. Undoes some pre-processing
        """

        def remove_underscore(m: re.Match) -> str:
            s: str = m.group(0)
            # All matches should only have a single underscore
            assert s.count("_") == 1, "Internal Error! Please open an issue"
            return s.replace("_", "")

        for repl_patt, n in string_properties.keyword_replacements.items():
            formatted_string, k = re.subn(
                repl_patt, remove_underscore, formatted_string
            )
            assert n == k, (
                "Internal Error! Did not reverse the same number of keywords that "
                "were replaced. Please open an issue"
            )

        for operator, n in string_properties.tilde_replacements.items():
            assert n == formatted_string.count(operator), (
                "Internal Error! Did not reverse the same number of tildes that "
                "were replaced. Please open an issue"
            )
            formatted_string = formatted_string.replace(operator, "~")

        return formatted_string


@dataclass
class JinjaTag:
    """
    A simple representation of a jinja tag.

    "verb" is one of {set, do, for, if, elif, else, test, macro, materialization, call}

    For example, "{%- set my_var=4 %}" is split into it parts:
    (opening_marker, verb, code, closing_marker) = ("{%-", "set", "my_var=4", "%}")
    """

    source_string: str
    opening_marker: str
    verb: str
    code: str
    closing_marker: str
    depth: Tuple[int, int]
    is_blackened: bool = False

    def __str__(self) -> str:
        if self.is_macro_like_def and self.is_blackened:
            self._remove_trailing_comma()

        if self.is_multiline_tag and self.is_blackened:
            return self._multiline_str()
        elif self.is_multiline_tag:
            return self.source_string
        else:
            return self._basic_str()

    @property
    def is_multiline_tag(self) -> bool:
        return "\n" in self.code

    @property
    def is_macro_like_def(self) -> bool:
        return "%" in self.opening_marker and any(
            [
                self.verb.strip() == v
                for v in ["macro", "test", "materialization", "call"]
            ]
        )

    def _multiline_str(self) -> str:
        """
        if the formatted code is on multiple lines, we want the code indented
        four spaces past the opening and closing markers. The opening marker
        will already be indented to the proper depth (because of the Line).
        """
        indent = " " * 4 * (self.depth[0] + self.depth[1])
        no_indent_lines = self._find_multiline_python_str_lines()
        code_lines = iter(self.code.splitlines(keepends=False))

        if self.verb:
            lines = [f"{self.opening_marker} {self.verb}{next(code_lines)}"]
            extra_indent = ""
        else:
            lines = [f"{self.opening_marker}"]
            extra_indent = " " * 4

        for i, code_line in enumerate(code_lines, start=1 if self.verb else 0):
            line_indent = "" if i in no_indent_lines else f"{indent}{extra_indent}"
            lines.append(f"{line_indent}{code_line}")

        if self.verb:
            lines[-1] = f"{indent}{lines[-1].lstrip()} {self.closing_marker}"
        else:
            lines.append(f"{indent}{self.closing_marker}")

        return "\n".join(lines)

    def _basic_str(self) -> str:
        return f"{self.opening_marker} {self.verb}{self.code} {self.closing_marker}"

    def _find_multiline_python_str_lines(self) -> MutableSet[int]:
        """
        Return a set line numbers that correspond with the lines
        of a triple-quoted multiline string (except for the first line
        of each string). These are lines that should never have
        their indentation adjusted
        """
        try:
            tree = ast.parse(self.code, mode="exec")
        except SyntaxError:
            # this jinja isn't quite python, so give up here.
            return set()

        line_indicies: MutableSet[int] = set()
        raw_lines = self.code.splitlines()
        for node in ast.walk(tree):
            if (
                isinstance(node, ast.Constant)
                and isinstance(node.value, str)
                and node.end_lineno is not None
                and node.end_lineno > node.lineno
                and "\n" in node.value
                and (
                    '"""' in raw_lines[node.lineno - 1]
                    or "'''" in raw_lines[node.lineno - 1]
                )
                and (
                    '"""' in raw_lines[node.end_lineno - 1]
                    or "'''" in raw_lines[node.end_lineno - 1]
                )
            ):
                line_indicies |= set(range(node.lineno, node.end_lineno))

        return line_indicies

    def _remove_trailing_comma(self) -> None:
        """
        dbt Jinja doesn't allow trailing commas in macro definitions. Mutates
        self.code so that it doesn't contain a trailing comma inside parentheses.
        """
        trailing_comma_prog = re.compile(r",\s*\)")
        source_string = self.code
        trailing_comma_match = trailing_comma_prog.search(source_string)
        if trailing_comma_match:
            idx = trailing_comma_match.span()[0]
            processed_string = f"{source_string[:idx]}{source_string[idx+1:]}"
            self.code = processed_string

    @classmethod
    def from_string(cls, source_string: str, depth: Tuple[int, int]) -> "JinjaTag":
        """
        Takes a jinja statement or expression as a string and returns
        a JinjaTag object (basically a tuple of its parts).
        """
        opening_marker_len = 3 if source_string[2] == "-" else 2
        opening_marker = source_string[:opening_marker_len]
        closing_marker_len = 3 if source_string[-3] == "-" else 2
        closing_marker = source_string[-closing_marker_len:]

        verb_pattern = (
            r"\s*(set|do|for|if|elif|else|test|macro|materialization|call)\s+"
        )
        verb_program = re.compile(verb_pattern, re.DOTALL | re.IGNORECASE)
        verb_match = verb_program.match(source_string[opening_marker_len:])
        if verb_match:
            verb_pos = opening_marker_len + verb_match.span(1)[0]
            verb_epos = opening_marker_len + verb_match.span(1)[1]
            verb = source_string[verb_pos:verb_epos].lower()
        else:
            verb = ""
        verb_len = verb_match.span(0)[1] - verb_match.span(0)[0] if verb_match else 0

        code_pos = opening_marker_len + verb_len
        code = source_string[code_pos:-closing_marker_len].strip()

        if verb and code:
            verb = f"{verb} "

        return JinjaTag(
            source_string, opening_marker, verb, code, closing_marker, depth
        )

    def max_code_length(self, max_length: int) -> int:
        """
        For a tag with max_length remaining characters on the line, return the
        max length that the code inside the curlies (exc. the verb) can occupy.
        """
        return (
            max_length
            - len(self.opening_marker)
            - len(self.verb)
            - len(self.closing_marker)
            - 2
        )


@dataclass
class JinjaFormatter:
    """
    Provides a simple interface, format_line, to format all jinja tags
    on a Line, using black if it is installed
    """

    mode: Mode
    code_formatter: BlackWrapper = field(default_factory=lambda: BlackWrapper())

    def __post_init__(self) -> None:
        self.use_black = (
            self.code_formatter.black is not None and not self.mode.no_jinjafmt
        )
        self.node_manager = NodeManager(self.mode.dialect.case_sensitive_names)

    def format_line(self, line: Line) -> List[Line]:
        """
        Format each jinja tag in a line, in turn. If a node was made multiline,
        split before the node (unless it is already the first node on that line)
        """
        line_length = self.mode.line_length
        if not line.formatting_disabled and line.contains_jinja:
            running_length = len(line.prefix)
            for i, node in enumerate(line.nodes):
                is_blackened = self._format_jinja_node(
                    node, max_length=line_length - running_length
                )
                if i > 0 and is_blackened and node.is_multiline_jinja:
                    # if black turned a single-line jinja expression
                    # into a multiline one, we need to split before
                    # this node (since on the first pass the splitter
                    # would have split before this node if it had been
                    # a multiline node)
                    splitter = LineSplitter(self.node_manager)
                    return list(
                        chain(
                            *[
                                self.format_line(new_line)
                                for new_line, _ in [
                                    splitter.split_at_index(line, 0, i, line.comments),
                                    splitter.split_at_index(
                                        line, i, -1, [], no_tail=True
                                    ),
                                ]
                            ]
                        )
                    )
                else:
                    running_length += len(node) - (len(node.prefix) if i == 0 else 0)
            else:
                return [line]
        else:
            return [line]

    def _format_jinja_node(self, node: Node, max_length: int) -> bool:
        """
        Format a single jinja tag. No-ops for nodes that
        are not jinja. Returns True if the node was blackened
        """
        if node.is_jinja:
            tag = JinjaTag.from_string(node.value, node.depth)

            if tag.code and self.use_black:
                tag.code, tag.is_blackened = self.code_formatter.format_string(
                    tag.code,
                    max_length=tag.max_code_length(max_length),
                )

            node.value = str(tag)

            return tag.is_blackened

        else:
            return False