pbi-parsers 0.7.8__py3-none-any.whl

pbi_parsers/dax/utils.py

@@ -0,0 +1,120 @@
+from typing import TypeVar
+
+import jinja2
+from colorama import Fore, Style
+
+from .exprs import Expression
+from .tokens import Token
+
+T = TypeVar("T", bound=Expression)
+
+
+def find_all(ast: Expression, class_type: type[T] | tuple[type[T], ...]) -> list[T]:
+    """Find all instances of a specific class type in the AST."""
+    ret = []
+    for child in ast.children():
+        if isinstance(child, class_type):
+            ret.append(child)
+        ret.extend(find_all(child, class_type))
+    return ret
+
+
+CONSOLE = jinja2.Template("""
+{%- for i, section_line in enumerate(lines) -%}
+{%- if i in highlights %}
+{{ Style.BRIGHT }}{{ Fore.CYAN }}{{ i }} |{{ Style.RESET_ALL }} {{ section_line }}
+{{ " " * (highlights[i][0]) }}{{ Style.BRIGHT }}{{ Fore.YELLOW }}{{ "^" * (highlights[i][1] - highlights[i][0]) }}{{ Style.RESET_ALL }}
+{%- elif i >= section_boundary_lines[0] and i <= section_boundary_lines[1] %}
+{{ i }} | {{ section_line }}
+{%- endif %}
+{%- endfor %}
+""")
+HTML = jinja2.Template("""
+<div>
+{% for i, section_line in enumerate(section_lines) %}
+    <span class="{{ "" if i == 0 or i == section_lines|length - 1 else "highlighted" }}">{{ starting_line + i }} |</span>
+    <span>{{ section_line }}</span>
+{% endfor %}
+<div>
+""")  # noqa: E501
+
+
+class Context:
+    position: tuple[int, int]
+    full_text: str
+
+    def __init__(self, position: tuple[int, int], full_text: str) -> None:
+        self.position = position
+        self.full_text = full_text
+
+    def __repr__(self) -> str:
+        return self.to_console()
+
+    @staticmethod
+    def _get_highlighted_text(
+        lines: list[str],
+        position: tuple[int, int],
+    ) -> dict[int, tuple[int, int]]:
+        highlight_line_dict: dict[int, tuple[int, int]] = {}
+
+        remaining_start, remaining_end = position
+        for i, line in enumerate(lines):
+            if len(line) > remaining_start and remaining_end > 0:
+                buffer = len(str(i)) + 3
+                highlight_line_dict[i] = (
+                    buffer + remaining_start,
+                    buffer + min(remaining_end, len(line)),
+                )
+            remaining_start -= len(line) + 1  # +1 for the newline character
+            remaining_end -= len(line) + 1
+        return highlight_line_dict
+
+    def to_console(self, context_lines: int = 2) -> str:
+        """Render the context for console output."""
+        lines = self.full_text.split("\n")
+        starting_line = self.full_text.count("\n", 0, self.position[0])
+        final_line = self.full_text.count("\n", 0, self.position[1])
+        section_boundary_lines = (max(starting_line - context_lines, 0), min(final_line + context_lines, len(lines)))
+        highlights = self._get_highlighted_text(lines, self.position)
+        return CONSOLE.render(
+            lines=lines,
+            section_boundary_lines=section_boundary_lines,
+            highlights=highlights,
+            enumerate=enumerate,
+            Style=Style,
+            Fore=Fore,
+        )
+
+    def to_dict(self) -> dict[str, str | tuple[int, int]]:
+        """Convert the context to a dictionary."""
+        return {
+            "position": self.position,
+            "full_text": self.full_text,
+        }
+
+    def to_html(self) -> str:
+        """Render the context for console output."""
+        lines = self.full_text.split("\n")
+        starting_line = self.full_text.count("\n", 0, self.position[0]) + 1
+        final_line = self.full_text.count("\n", 0, self.position[1]) + 1
+
+        section_lines = lines[starting_line - 2 : final_line + 1]
+        return HTML.render(
+            section_lines=section_lines,
+            enumerate=enumerate,
+            starting_line=starting_line,
+            Style=Style,
+            Fore=Fore,
+        )
+
+
+def highlight_section(node: Expression | Token | list[Token] | list[Expression]) -> Context:
+    if isinstance(node, list):
+        position = (node[0].position()[0], node[-1].position()[1])
+        first_node = node[0]
+        full_text = first_node.text_slice.full_text if isinstance(first_node, Token) else first_node.full_text()
+        return Context(position, full_text)
+
+    position = node.position()
+    full_text = node.text_slice.full_text if isinstance(node, Token) else node.full_text()
+    return Context(position, full_text)

pbi_parsers/pq/__init__.py

@@ -0,0 +1,17 @@
+from .exprs import Expression
+from .formatter import Formatter
+from .lexer import Lexer
+from .main import format_expression, to_ast
+from .parser import Parser
+from .tokens import Token, TokenType
+
+__all__ = [
+    "Expression",
+    "Formatter",
+    "Lexer",
+    "Parser",
+    "Token",
+    "TokenType",
+    "format_expression",
+    "to_ast",
+]

pbi_parsers/pq/exprs/__init__.py

@@ -0,0 +1,98 @@
+from typing import TYPE_CHECKING
+
+from ._base import Expression
+from .add_sub import AddSubExpression
+from .add_sub_unary import AddSubUnaryExpression
+from .and_or_expr import AndOrExpression
+from .array import ArrayExpression
+from .arrow import ArrowExpression
+from .column import ColumnExpression
+from .comparison import ComparisonExpression
+from .concatenation import ConcatenationExpression
+from .div_mul import DivMulExpression
+from .each import EachExpression
+from .ellipsis_expr import EllipsisExpression
+from .function import FunctionExpression
+from .identifier import BracketedIdentifierExpression, IdentifierExpression
+from .if_expr import IfExpression
+from .is_expr import IsExpression
+from .keyword import KeywordExpression
+from .literal_number import LiteralNumberExpression
+from .literal_string import LiteralStringExpression
+from .meta import MetaExpression
+from .negation import NegationExpression
+from .not_expr import NotExpression
+from .parens import ParenthesesExpression
+from .record import RecordExpression
+from .row import RowExpression
+from .row_index import RowIndexExpression
+from .statement import StatementExpression
+from .try_expr import TryExpression
+from .type_expr import TypingExpression
+from .variable import VariableExpression
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+"""
+operators > all
+
+operators:
+if, comparison, concatenation, add/sub, div/mul, meta, add/sub, unary
+
+variable > identifier
+row > identifier
+function > identifier
+typing > keyword
+arrow > parentheses
+"""
+EXPRESSION_HIERARCHY: tuple[type[Expression], ...] = (
+    IfExpression,
+    AndOrExpression,
+    ComparisonExpression,
+    ConcatenationExpression,
+    AddSubExpression,
+    DivMulExpression,
+    IsExpression,
+    MetaExpression,
+    NegationExpression,
+    AddSubUnaryExpression,
+    RowIndexExpression,
+    # above are left-associative operators
+    NotExpression,
+    EllipsisExpression,
+    ArrowExpression,
+    TryExpression,
+    ParenthesesExpression,
+    StatementExpression,
+    ColumnExpression,
+    EachExpression,
+    ArrayExpression,
+    FunctionExpression,
+    VariableExpression,
+    RowExpression,
+    TypingExpression,
+    KeywordExpression,
+    IdentifierExpression,
+    RecordExpression,
+    BracketedIdentifierExpression,
+    LiteralStringExpression,
+    LiteralNumberExpression,
+)
+
+
+def any_expression_match(parser: "Parser", skip_first: int = 0) -> Expression | None:
+    """Matches any expression type.
+
+    This is a utility function to simplify the matching process in other expressions.
+    """
+    for expr in EXPRESSION_HIERARCHY[skip_first:]:
+        if match := expr.match(parser):
+            return match
+    return None
+
+
+__all__ = [
+    "Expression",
+    "any_expression_match",
+]

pbi_parsers/pq/exprs/_base.py

@@ -0,0 +1,33 @@
+from typing import TYPE_CHECKING
+
+from pbi_parsers.pq.tokens import TokenType
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+
+class Expression:
+    def pprint(self) -> str:
+        msg = "Subclasses should implement this method."
+        raise NotImplementedError(msg)
+
+    @classmethod
+    def match(cls, parser: "Parser") -> "Expression | None":
+        """Attempt to match the current tokens to this expression type.
+
+        Returns an instance of the expression if matched, otherwise None.
+        """
+        msg = "Subclasses should implement this method."
+        raise NotImplementedError(msg)
+
+    @staticmethod
+    def match_tokens(parser: "Parser", match_tokens: list[TokenType]) -> bool:
+        return all(parser.peek(i).tok_type == token_type for i, token_type in enumerate(match_tokens))
+
+    def __repr__(self) -> str:
+        return self.pprint()
+
+    def children(self) -> list["Expression"]:
+        """Returns a list of child expressions."""
+        msg = "This method should be implemented by subclasses."
+        raise NotImplementedError(msg)

pbi_parsers/pq/exprs/_utils.py

@@ -0,0 +1,31 @@
+from collections.abc import Callable
+from typing import ParamSpec, TypeVar
+
+from pbi_parsers.pq.parser import Parser
+
+P = ParamSpec("P")  # Represents the parameters of the decorated function
+R = TypeVar("R")  # Represents the return type of the decorated function
+
+
+def lexer_reset(func: Callable[P, R]) -> Callable[P, R]:
+    def lexer_reset_inner(*args: P.args, **kwargs: P.kwargs) -> R:
+        parser = args[1]
+        if not isinstance(parser, Parser):
+            msg = f"Expected the second argument to be a Parser instance, got {type(parser)}"
+            raise TypeError(msg)
+        idx = parser.index
+
+        # Speed up of a bazillion
+        cached_val, cached_index = parser.cache.get((idx, id(func)), (None, -1))
+        if cached_val is not None:
+            parser.index = cached_index
+            return cached_val
+
+        ret = func(*args, **kwargs)
+
+        parser.cache[idx, id(func)] = (ret, parser.index)
+        if ret is None:
+            parser.index = idx
+        return ret
+
+    return lexer_reset_inner

pbi_parsers/pq/exprs/add_sub.py

@@ -0,0 +1,59 @@
+import textwrap
+from typing import TYPE_CHECKING
+
+from pbi_parsers.pq.tokens import Token, TokenType
+
+from ._base import Expression
+from ._utils import lexer_reset
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+
+class AddSubExpression(Expression):
+    """Represents an addition or subtraction expression."""
+
+    operator: Token
+    left: Expression
+    right: Expression
+
+    def __init__(self, operator: Token, left: Expression, right: Expression) -> None:
+        self.operator = operator
+        self.left = left
+        self.right = right
+
+    @classmethod
+    @lexer_reset
+    def match(cls, parser: "Parser") -> "AddSubExpression | None":
+        from . import EXPRESSION_HIERARCHY, any_expression_match  # noqa: PLC0415
+
+        skip_index = EXPRESSION_HIERARCHY.index(AddSubExpression)
+
+        left_term = any_expression_match(parser=parser, skip_first=skip_index + 1)
+        operator = parser.consume()
+
+        if not left_term:
+            return None
+        if operator.tok_type not in {TokenType.PLUS_SIGN, TokenType.MINUS_SIGN}:
+            return None
+
+        right_term = any_expression_match(parser=parser, skip_first=skip_index)
+        if right_term is None:
+            msg = f"Expected a right term after operator {operator.text}, found: {parser.peek()}"
+            raise ValueError(msg)
+        return AddSubExpression(operator=operator, left=left_term, right=right_term)
+
+    def pprint(self) -> str:
+        op_str = "Add" if self.operator.text == "+" else "Sub"
+        left_str = textwrap.indent(self.left.pprint(), " " * 10).lstrip()
+        right_str = textwrap.indent(self.right.pprint(), " " * 10).lstrip()
+        return f"""
+{op_str} (
+    operator: {self.operator.text},
+    left: {left_str},
+    right: {right_str}
+)""".strip()
+
+    def children(self) -> list[Expression]:
+        """Returns a list of child expressions."""
+        return [self.left, self.right]

pbi_parsers/pq/exprs/add_sub_unary.py

@@ -0,0 +1,57 @@
+import textwrap
+from typing import TYPE_CHECKING
+
+from pbi_parsers.pq.tokens import Token, TokenType
+
+from ._base import Expression
+from ._utils import lexer_reset
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+
+class AddSubUnaryExpression(Expression):
+    """Represents an addition or subtraction expression."""
+
+    operator: Token
+    number: Expression
+
+    def __init__(self, operator: Token, number: Expression) -> None:
+        self.operator = operator
+        self.number = number
+
+    @classmethod
+    @lexer_reset
+    def match(cls, parser: "Parser") -> "AddSubUnaryExpression | None":
+        from . import EXPRESSION_HIERARCHY, any_expression_match  # noqa: PLC0415
+
+        skip_index = EXPRESSION_HIERARCHY.index(
+            AddSubUnaryExpression,
+        )  # intentionally inclusive of self to allow +-++- chains
+
+        operator = parser.consume()
+
+        if operator.tok_type not in {TokenType.PLUS_SIGN, TokenType.MINUS_SIGN}:
+            return None
+
+        # Handle chained +-++-+ prefixes
+        number: Expression | None = any_expression_match(
+            parser=parser,
+            skip_first=skip_index,
+        )
+        if number is None:
+            msg = f"Expected a number after operator {operator.text}, found: {parser.peek()}"
+            raise ValueError(msg)
+        return AddSubUnaryExpression(operator=operator, number=number)
+
+    def pprint(self) -> str:
+        number = textwrap.indent(self.number.pprint(), " " * 12).lstrip()
+        return f"""
+Number (
+    sign: {self.operator.text},
+    number: {number},
+)""".strip()
+
+    def children(self) -> list[Expression]:
+        """Returns a list of child expressions."""
+        return [self.number]

pbi_parsers/pq/exprs/and_or_expr.py

@@ -0,0 +1,60 @@
+import textwrap
+from typing import TYPE_CHECKING
+
+from pbi_parsers.pq.tokens import Token, TokenType
+
+from ._base import Expression
+from ._utils import lexer_reset
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+
+class AndOrExpression(Expression):
+    """Represents an AND or OR expression."""
+
+    operator: Token
+    left: Expression
+    right: Expression
+
+    def __init__(self, operator: Token, left: Expression, right: Expression) -> None:
+        self.operator = operator
+        self.left = left
+        self.right = right
+
+    @classmethod
+    @lexer_reset
+    def match(cls, parser: "Parser") -> "AndOrExpression | None":
+        from . import EXPRESSION_HIERARCHY, any_expression_match  # noqa: PLC0415
+
+        skip_index = EXPRESSION_HIERARCHY.index(AndOrExpression)
+
+        left_term = any_expression_match(parser=parser, skip_first=skip_index + 1)
+        operator = parser.consume()
+
+        if not left_term:
+            return None
+        if operator.tok_type not in {
+            TokenType.AND,
+            TokenType.OR,
+        }:
+            return None
+
+        right_term = any_expression_match(parser=parser, skip_first=skip_index)
+        if right_term is None:
+            msg = f"Expected a right term after operator {operator.text}, found: {parser.peek()}"
+            raise ValueError(msg)
+        return AndOrExpression(operator=operator, left=left_term, right=right_term)
+
+    def pprint(self) -> str:
+        left_str = textwrap.indent(self.left.pprint(), " " * 10)[10:]
+        right_str = textwrap.indent(self.right.pprint(), " " * 10)[10:]
+        return f"""
+{self.operator.text} (
+    left: {left_str},
+    right: {right_str}
+)""".strip()
+
+    def children(self) -> list[Expression]:
+        """Returns a list of child expressions."""
+        return [self.left, self.right]

pbi_parsers/pq/exprs/array.py

@@ -0,0 +1,53 @@
+import textwrap
+from typing import TYPE_CHECKING
+
+from pbi_parsers.pq.tokens import TokenType
+
+from ._base import Expression
+from ._utils import lexer_reset
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+
+class ArrayExpression(Expression):
+    elements: list[Expression]
+
+    def __init__(self, elements: list[Expression]) -> None:
+        self.elements: list[Expression] = elements
+
+    def pprint(self) -> str:
+        elements = ",\n".join(element.pprint() for element in self.elements)
+        elements = textwrap.indent(elements, " " * 14)[14:]
+        return f"""
+Array (
+    elements: {elements}
+)        """.strip()
+
+    @classmethod
+    @lexer_reset
+    def match(cls, parser: "Parser") -> "ArrayExpression | None":
+        from . import any_expression_match  # noqa: PLC0415
+
+        if parser.consume().tok_type != TokenType.LEFT_CURLY_BRACE:
+            return None
+
+        elements: list[Expression] = []
+
+        while not cls.match_tokens(parser, [TokenType.RIGHT_CURLY_BRACE]):
+            # We gotta handle operators next :(
+            element = any_expression_match(parser)
+            if element is not None:
+                elements.append(element)
+            else:
+                msg = f"Unexpected token sequence: {parser.peek()}, {parser.index}"
+                raise ValueError(msg)
+
+            if not cls.match_tokens(parser, [TokenType.RIGHT_CURLY_BRACE]):
+                assert parser.consume().tok_type == TokenType.COMMA
+        _right_brace = parser.consume()
+        return ArrayExpression(elements=elements)
+
+    def children(self) -> list[Expression]:
+        """Returns a list of child expressions."""
+        return self.elements

pbi_parsers/pq/exprs/arrow.py

@@ -0,0 +1,50 @@
+import textwrap
+from typing import TYPE_CHECKING
+
+from pbi_parsers.pq.tokens import TokenType
+
+from ._base import Expression
+from ._utils import lexer_reset
+from .parens import ParenthesesExpression
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+
+class ArrowExpression(Expression):
+    inputs: ParenthesesExpression
+    function_body: Expression
+
+    def __init__(self, inputs: ParenthesesExpression, function_body: Expression) -> None:
+        self.inputs = inputs
+        self.function_body = function_body
+
+    def pprint(self) -> str:
+        inputs = textwrap.indent(self.inputs.pprint(), " " * 10)[10:]
+        function_body = textwrap.indent(self.function_body.pprint(), " " * 10)[10:]
+        return f"""
+Arrow (
+    inputs: {inputs},
+    function_body: {function_body}
+)""".strip()
+
+    @classmethod
+    @lexer_reset
+    def match(cls, parser: "Parser") -> "ArrowExpression | None":
+        from . import any_expression_match  # noqa: PLC0415
+
+        inputs = ParenthesesExpression.match(parser)
+        if inputs is None:
+            return None
+
+        if parser.consume().tok_type != TokenType.LAMBDA_ARROW:
+            return None
+        function_body = any_expression_match(parser)
+        if function_body is None:
+            return None
+
+        return ArrowExpression(inputs=inputs, function_body=function_body)
+
+    def children(self) -> list[Expression]:
+        """Returns a list of child expressions."""
+        return [self.inputs, self.function_body]

pbi_parsers/pq/exprs/column.py

@@ -0,0 +1,42 @@
+from typing import TYPE_CHECKING
+
+from pbi_parsers.pq.tokens import Token, TokenType
+
+from ._base import Expression
+from ._utils import lexer_reset
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+
+class ColumnExpression(Expression):
+    name: Token
+
+    def __init__(self, name: Token) -> None:
+        self.name = name
+
+    def pprint(self) -> str:
+        return f"Column ({self.name.text})"
+
+    @classmethod
+    @lexer_reset
+    def match(cls, parser: "Parser") -> "ColumnExpression | None":
+        if cls.match_tokens(
+            parser,
+            [
+                TokenType.LEFT_BRACKET,
+                TokenType.UNQUOTED_IDENTIFIER,
+                TokenType.RIGHT_BRACKET,
+            ],
+        ):
+            _l_bracket, name, _r_bracket = (
+                parser.consume(),
+                parser.consume(),
+                parser.consume(),
+            )
+            return ColumnExpression(name=name)
+        return None
+
+    def children(self) -> list[Expression]:  # noqa: PLR6301
+        """Returns a list of child expressions."""
+        return []

pbi_parsers/pq/exprs/comparison.py

@@ -0,0 +1,62 @@
+import textwrap
+from typing import TYPE_CHECKING
+
+from pbi_parsers.pq.tokens import Token, TokenType
+
+from ._base import Expression
+from ._utils import lexer_reset
+
+if TYPE_CHECKING:
+    from pbi_parsers.pq.parser import Parser
+
+
+class ComparisonExpression(Expression):
+    """Represents an multiplication or division expression."""
+
+    operator: Token
+    left: Expression
+    right: Expression
+
+    def __init__(self, operator: Token, left: Expression, right: Expression) -> None:
+        self.operator = operator
+        self.left = left
+        self.right = right
+
+    @classmethod
+    @lexer_reset
+    def match(cls, parser: "Parser") -> "ComparisonExpression | None":
+        from . import EXPRESSION_HIERARCHY, any_expression_match  # noqa: PLC0415
+
+        skip_index = EXPRESSION_HIERARCHY.index(ComparisonExpression)
+
+        left_term = any_expression_match(parser=parser, skip_first=skip_index + 1)
+        operator = parser.consume()
+
+        if not left_term:
+            return None
+        if operator.tok_type not in {
+            TokenType.EQUAL_SIGN,
+            TokenType.COMPARISON_OPERATOR,
+            TokenType.NOT_EQUAL_SIGN,
+        }:
+            return None
+
+        right_term = any_expression_match(parser=parser, skip_first=skip_index)
+        if right_term is None:
+            msg = f"Expected a right term after operator {operator.text}, found: {parser.peek()}"
+            raise ValueError(msg)
+        return ComparisonExpression(operator=operator, left=left_term, right=right_term)
+
+    def pprint(self) -> str:
+        left_str = textwrap.indent(self.left.pprint(), " " * 10)[10:]
+        right_str = textwrap.indent(self.right.pprint(), " " * 10)[10:]
+        return f"""
+Bool (
+    operator: {self.operator.text},
+    left: {left_str},
+    right: {right_str}
+)""".strip()
+
+    def children(self) -> list[Expression]:
+        """Returns a list of child expressions."""
+        return [self.left, self.right]