strukto-ai--mirage/python/mirage/commands/builtin/generic/zgrep.py

import gzip as gziplib
import re
from collections.abc import (AsyncIterator, Awaitable, Callable, Mapping,
                             Sequence)
from dataclasses import dataclass
from functools import partial

from mirage.accessor.base import Accessor
from mirage.commands.builtin.grep_helper import (build_pattern_str,
                                                 resolve_pattern)
from mirage.commands.builtin.utils.lines import split_lines
from mirage.commands.builtin.utils.stream import _read_stdin_async
from mirage.commands.spec import SPECS
from mirage.commands.spec.types import FlagView
from mirage.io.types import ByteSource, IOResult
from mirage.types import PathSpec


async def _read_plain(
    read_bytes: Callable[..., Awaitable[bytes]],
    accessor: Accessor,
    path: PathSpec,
    index: object = None,
) -> bytes:
    return await read_bytes(accessor, path)


def _zgrep_search(
    data: bytes,
    pattern: str,
    ignore_case: bool,
    invert: bool,
    count: bool,
    line_numbers: bool,
    filename: str | None,
    only_matching: bool,
    max_count: int | None,
) -> tuple[list[str], bool]:
    text = data.decode(errors="replace")
    lines = split_lines(text)
    flags = re.IGNORECASE if ignore_case else 0
    matched: list[tuple[int, str]] = []
    for idx, line in enumerate(lines, 1):
        if only_matching and not invert:
            hits = list(re.finditer(pattern, line, flags))
            if hits:
                for m in hits:
                    matched.append((idx, m.group()))
                    if max_count is not None and len(matched) >= max_count:
                        break
            elif invert:
                matched.append((idx, line))
        else:
            hit = bool(re.search(pattern, line, flags))
            if invert:
                hit = not hit
            if hit:
                matched.append((idx, line))
        if max_count is not None and len(matched) >= max_count:
            break
    if count:
        value = str(len(matched))
        if filename:
            value = f"{filename}:{value}"
        return [value], len(matched) > 0
    result: list[str] = []
    for idx, line in matched:
        prefix = ""
        if filename:
            prefix = filename + ":"
        if line_numbers:
            prefix += str(idx) + ":"
        result.append(prefix + line)
    return result, len(matched) > 0


def _files_only_match(data: bytes, pattern: str, ignore_case: bool,
                      invert: bool) -> bool:
    text = data.decode(errors="replace")
    flags = re.IGNORECASE if ignore_case else 0
    for line in split_lines(text):
        hit = bool(re.search(pattern, line, flags))
        if invert:
            hit = not hit
        if hit:
            return True
    return False


@dataclass(frozen=True, slots=True)
class ZgrepFlags:
    """Parsed zgrep flags; the complete set zgrep honors."""
    ignore_case: bool
    invert: bool
    count: bool
    files_only: bool
    line_numbers: bool
    fixed: bool
    force_filename: bool
    suppress_filename: bool
    only_matching: bool
    quiet: bool
    whole_word: bool
    max_count: int | None


def parse_flags(fl: FlagView, never_match: bool) -> ZgrepFlags:
    """Convert the raw flag bag into ZgrepFlags, the only string-keyed reads.

    Args:
        fl (FlagView): spec-validated view over the raw flag kwargs.
        never_match (bool): zero-pattern sentinel from resolve_pattern; it is
            a regex, so it suppresses -F.
    """
    return ZgrepFlags(
        ignore_case=fl.bool("i"),
        invert=fl.bool("v"),
        count=fl.bool("c"),
        files_only=fl.bool("args_l"),
        line_numbers=fl.bool("n"),
        fixed=fl.bool("F") and not never_match,
        force_filename=fl.bool("H"),
        suppress_filename=fl.bool("h"),
        only_matching=fl.bool("o"),
        quiet=fl.bool("q"),
        whole_word=fl.bool("w"),
        max_count=fl.int("m"),
    )


async def zgrep(
    paths: list[PathSpec],
    texts: Sequence[str] = (),
    flags: Mapping[str, object] | None = None,
    *,
    read_bytes: Callable[..., Awaitable[bytes]],
    accessor: Accessor | None = None,
    stdin: AsyncIterator[bytes] | bytes | None = None,
    index: object = None,
) -> tuple[ByteSource | None, IOResult]:
    fl = FlagView(flags, spec=SPECS["zgrep"])
    pattern, never_match = await resolve_pattern(
        texts, fl, partial(_read_plain, read_bytes), accessor, index,
        "zgrep: usage: zgrep [flags] pattern [path]")
    f = parse_flags(fl, never_match)
    compiled = build_pattern_str(pattern, f.fixed, f.whole_word)
    multi = len(paths) > 1
    show_filename = f.force_filename or (multi and not f.suppress_filename)
    any_match = False
    all_results: list[str] = []

    if paths:
        for p in paths:
            raw = await read_bytes(accessor, p)
            data = gziplib.decompress(raw)
            fname = p.virtual if show_filename else None
            if f.files_only:
                if _files_only_match(data, compiled, f.ignore_case, f.invert):
                    all_results.append(p.virtual)
                    any_match = True
            else:
                result, had_match = _zgrep_search(data, compiled,
                                                  f.ignore_case, f.invert,
                                                  f.count, f.line_numbers,
                                                  fname, f.only_matching,
                                                  f.max_count)
                if had_match:
                    any_match = True
                all_results.extend(result)
    else:
        raw = await _read_stdin_async(stdin)
        data = gziplib.decompress(raw) if raw else b""
        if f.files_only:
            if _files_only_match(data, compiled, f.ignore_case, f.invert):
                all_results.append("(standard input)")
                any_match = True
        else:
            result, had_match = _zgrep_search(data, compiled, f.ignore_case,
                                              f.invert, f.count,
                                              f.line_numbers, None,
                                              f.only_matching, f.max_count)
            if had_match:
                any_match = True
            all_results.extend(result)

    if f.quiet:
        return None, IOResult(exit_code=0 if any_match else 1)
    exit_code = 0 if any_match else 1
    if not all_results:
        return None, IOResult(exit_code=exit_code)
    return ("\n".join(all_results) +
            "\n").encode(), IOResult(exit_code=exit_code)


__all__ = ["zgrep"]