tgit/app/logs/grep.py

"""`grep` for logs: parse a grep-like expression, run it with the `regex` module, show it back as a command.

Nothing here executes anything: `shlex.split` only tokenizes, the pattern goes to `regex` with a per-line timeout
and an overall time budget, so `(a+)+$` on a long line costs at most the budget, not the bot.
Supported: PATTERN, -i, -v, -E, -c, -C N (also -C2, -iv, `--`, a leading `grep`). Anything else is an error.
"""

import shlex
import time
from dataclasses import asdict, dataclass, field
from enum import StrEnum
from typing import Any

import regex

MAX_PATTERN = 200
MAX_CONTEXT = 10
MAX_LINES = 500  # output lines, context and `--` separators included
LINE_TIMEOUT = 0.05  # seconds per line
BUDGET = 3.0  # seconds per search; on overrun the result is partial and says so

_BRE_SPECIAL = set("+?|(){}")


class GrepErrorCode(StrEnum):
    EMPTY = "empty"
    QUOTES = "quotes"
    FLAG = "flag"
    CONTEXT = "context"
    EXTRA = "extra"
    TOO_LONG = "too_long"
    REGEX = "regex"


class GrepError(ValueError):
    def __init__(self, code: GrepErrorCode, arg: str = "") -> None:
        super().__init__(f"{code}: {arg}" if arg else code.value)
        self.code = code
        self.arg = arg


@dataclass(frozen=True)
class GrepQuery:
    pattern: str
    ignore_case: bool = False
    invert: bool = False
    extended: bool = False
    count: bool = False
    context: int = 0

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    def compile(self) -> "regex.Pattern[str]":
        source = self.pattern if self.extended else bre_to_ere(self.pattern)
        try:
            return regex.compile(source, regex.IGNORECASE if self.ignore_case else 0)
        except regex.error as e:
            raise GrepError(GrepErrorCode.REGEX, str(e)) from e


def bre_to_ere(pattern: str) -> str:
    """Basic grep regex → Python: without -E `+ ? | ( ) { }` are literals and `\\+` etc. are the operators."""
    out: list[str] = []
    i, in_class = 0, False
    while i < len(pattern):
        ch = pattern[i]
        if in_class:
            out.append(ch)
            if ch == "]" and not (out[-2:-1] == ["["] or out[-3:-1] == ["[", "^"]):
                in_class = False
        elif ch == "[":
            in_class = True
            out.append(ch)
        elif ch == "\\" and i + 1 < len(pattern):
            nxt = pattern[i + 1]
            out.append(nxt if nxt in _BRE_SPECIAL else ch + nxt)
            i += 1
        elif ch in _BRE_SPECIAL:
            out.append("\\" + ch)
        else:
            out.append(ch)
        i += 1
    return "".join(out)


def _context(value: str) -> int:
    if not value.isdigit() or int(value) > MAX_CONTEXT:
        raise GrepError(GrepErrorCode.CONTEXT, value)
    return int(value)


def parse_grep(expr: str) -> GrepQuery:
    try:
        tokens = shlex.split(expr)
    except ValueError as e:
        raise GrepError(GrepErrorCode.QUOTES) from e
    if tokens[:1] == ["grep"]:
        tokens = tokens[1:]
    opts: dict[str, Any] = {}
    patterns: list[str] = []
    i = 0
    while i < len(tokens):
        tok = tokens[i]
        if patterns or tok == "--" or not tok.startswith("-") or tok == "-":
            if tok == "--" and not patterns:
                patterns.extend(tokens[i + 1 :])
                break
            patterns.append(tok)
        elif tok.startswith("-C"):
            if len(tok) > 2:
                opts["context"] = _context(tok[2:])
            elif i + 1 < len(tokens):
                opts["context"] = _context(tokens[i + 1])
                i += 1
            else:
                raise GrepError(GrepErrorCode.CONTEXT, "")
        else:
            for j, flag in enumerate(tok[1:], start=1):
                match flag:
                    case "i":
                        opts["ignore_case"] = True
                    case "v":
                        opts["invert"] = True
                    case "E":
                        opts["extended"] = True
                    case "c":
                        opts["count"] = True
                    case "C" if tok[j + 1 :].isdigit():
                        opts["context"] = _context(tok[j + 1 :])
                        break
                    case _:
                        raise GrepError(GrepErrorCode.FLAG, f"-{flag}")
        i += 1
    if not patterns or not patterns[0]:
        raise GrepError(GrepErrorCode.EMPTY)
    if len(patterns) > 1:
        raise GrepError(GrepErrorCode.EXTRA, patterns[1])
    if len(patterns[0]) > MAX_PATTERN:
        raise GrepError(GrepErrorCode.TOO_LONG, str(len(patterns[0])))
    query = GrepQuery(pattern=patterns[0], **opts)
    query.compile()  # syntax errors now, not halfway through the logs
    return query


def to_command(query: GrepQuery) -> str:
    """The query as a shell-looking line, for the user to see what was searched (SPEC L-15)."""
    flags = "".join(
        f for f, on in (("i", query.ignore_case), ("v", query.invert), ("E", query.extended), ("c", query.count)) if on
    )
    parts = ["grep"]
    if flags:
        parts.append(f"-{flags}")
    if query.context:
        parts += ["-C", str(query.context)]
    if query.pattern.startswith("-"):
        parts.append("--")
    parts.append(shlex.quote(query.pattern))
    return " ".join(parts)


@dataclass
class GrepResult:
    lines: list[str] = field(default_factory=list)
    matched: int = 0
    truncated: bool = False  # more output than MAX_LINES
    partial: bool = False  # time budget ran out: not every line was searched
    skipped: int = 0  # lines where the pattern hit the per-line timeout


def run(query: GrepQuery, text: str, max_lines: int = MAX_LINES, budget: float = BUDGET) -> GrepResult:
    pattern = query.compile()
    lines = text.splitlines()
    result = GrepResult()
    deadline = time.monotonic() + budget
    last_emitted = -1  # index of the last line already in the output (for context merging)
    for n, line in enumerate(lines):
        if time.monotonic() > deadline:
            result.partial = True
            break
        try:
            hit = pattern.search(line, timeout=LINE_TIMEOUT) is not None
        except TimeoutError:
            result.skipped += 1
            continue
        if hit == query.invert:
            continue
        result.matched += 1
        if query.count or result.truncated:
            continue
        start = max(n - query.context, last_emitted + 1)
        if query.context and result.lines and start > last_emitted + 1:
            result.lines.append("--")
        end = min(n + query.context, len(lines) - 1)
        for k in range(start, end + 1):
            if len(result.lines) >= max_lines:
                result.truncated = True
                break
            result.lines.append(lines[k])
            last_emitted = k
    return result