"""`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