codekingpro/portable-devtools

1"""Internal representation of a structured query language."""2 3from __future__ import annotations4 5from abc import ABC, abstractmethod6from enum import Enum7from typing import TYPE_CHECKING, Any8 9from pydantic import BaseModel10 11if TYPE_CHECKING:12    from collections.abc import Sequence13 14 15class Visitor(ABC):16    """Defines interface for IR translation using a visitor pattern."""17 18    allowed_comparators: Sequence[Comparator] | None = None19    """Allowed comparators for the visitor."""20 21    allowed_operators: Sequence[Operator] | None = None22    """Allowed operators for the visitor."""23 24    def _validate_func(self, func: Operator | Comparator) -> None:25        if (26            isinstance(func, Operator)27            and self.allowed_operators is not None28            and func not in self.allowed_operators29        ):30            msg = (31                f"Received disallowed operator {func}. Allowed "32                f"comparators are {self.allowed_operators}"33            )34            raise ValueError(msg)35        if (36            isinstance(func, Comparator)37            and self.allowed_comparators is not None38            and func not in self.allowed_comparators39        ):40            msg = (41                f"Received disallowed comparator {func}. Allowed "42                f"comparators are {self.allowed_comparators}"43            )44            raise ValueError(msg)45 46    @abstractmethod47    def visit_operation(self, operation: Operation) -> Any:48        """Translate an Operation.49 50        Args:51            operation: Operation to translate.52        """53 54    @abstractmethod55    def visit_comparison(self, comparison: Comparison) -> Any:56        """Translate a Comparison.57 58        Args:59            comparison: Comparison to translate.60        """61 62    @abstractmethod63    def visit_structured_query(self, structured_query: StructuredQuery) -> Any:64        """Translate a StructuredQuery.65 66        Args:67            structured_query: StructuredQuery to translate.68        """69 70 71def _to_snake_case(name: str) -> str:72    """Convert a name into snake_case."""73    snake_case = ""74    for i, char in enumerate(name):75        if char.isupper() and i != 0:76            snake_case += "_" + char.lower()77        else:78            snake_case += char.lower()79    return snake_case80 81 82class Expr(BaseModel):83    """Base class for all expressions."""84 85    def accept(self, visitor: Visitor) -> Any:86        """Accept a visitor.87 88        Args:89            visitor: visitor to accept.90 91        Returns:92            result of visiting.93        """94        return getattr(visitor, f"visit_{_to_snake_case(self.__class__.__name__)}")(95            self96        )97 98 99class Operator(str, Enum):100    """Enumerator of the operations."""101 102    AND = "and"103    OR = "or"104    NOT = "not"105 106 107class Comparator(str, Enum):108    """Enumerator of the comparison operators."""109 110    EQ = "eq"111    NE = "ne"112    GT = "gt"113    GTE = "gte"114    LT = "lt"115    LTE = "lte"116    CONTAIN = "contain"117    LIKE = "like"118    IN = "in"119    NIN = "nin"120 121 122class FilterDirective(Expr, ABC):123    """Filtering expression."""124 125 126class Comparison(FilterDirective):127    """Comparison to a value."""128 129    comparator: Comparator130    """The comparator to use."""131 132    attribute: str133    """The attribute to compare."""134 135    value: Any136    """The value to compare to."""137 138    def __init__(139        self, comparator: Comparator, attribute: str, value: Any, **kwargs: Any140    ) -> None:141        """Create a Comparison.142 143        Args:144            comparator: The comparator to use.145            attribute: The attribute to compare.146            value: The value to compare to.147        """148        # super exists from BaseModel149        super().__init__(150            comparator=comparator, attribute=attribute, value=value, **kwargs151        )152 153 154class Operation(FilterDirective):155    """Logical operation over other directives."""156 157    operator: Operator158    """The operator to use."""159 160    arguments: list[FilterDirective]161    """The arguments to the operator."""162 163    def __init__(164        self, operator: Operator, arguments: list[FilterDirective], **kwargs: Any165    ) -> None:166        """Create an Operation.167 168        Args:169            operator: The operator to use.170            arguments: The arguments to the operator.171        """172        # super exists from BaseModel173        super().__init__(operator=operator, arguments=arguments, **kwargs)174 175 176class StructuredQuery(Expr):177    """Structured query."""178 179    query: str180    """Query string."""181 182    filter: FilterDirective | None183    """Filtering expression."""184 185    limit: int | None186    """Limit on the number of results."""187 188    def __init__(189        self,190        query: str,191        filter: FilterDirective | None,  # noqa: A002192        limit: int | None = None,193        **kwargs: Any,194    ) -> None:195        """Create a StructuredQuery.196 197        Args:198            query: The query string.199            filter: The filtering expression.200            limit: The limit on the number of results.201        """202        # super exists from BaseModel203        super().__init__(query=query, filter=filter, limit=limit, **kwargs)204