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

import re
from collections.abc import AsyncIterator, Awaitable, Callable

from mirage.accessor.base import Accessor
from mirage.commands.builtin.utils.lines import split_lines
from mirage.commands.builtin.utils.stream import _read_stdin_async
from mirage.io.types import ByteSource, IOResult
from mirage.types import PathSpec


def _split_by_patterns(
    lines: list[str],
    patterns: list[str],
) -> list[list[str]]:
    parts: list[list[str]] = []
    current_start = 0
    for pat in patterns:
        if pat.startswith("/") and pat.endswith("/"):
            regex = pat[1:-1]
            for idx in range(current_start, len(lines)):
                if re.search(regex, lines[idx]):
                    parts.append(lines[current_start:idx])
                    current_start = idx
                    break
        else:
            line_num = int(pat)
            split_at = line_num - 1
            if split_at > current_start:
                parts.append(lines[current_start:split_at])
                current_start = split_at
    if current_start < len(lines):
        parts.append(lines[current_start:])
    return parts


async def csplit(
    paths: list[PathSpec],
    patterns: tuple[str, ...],
    *,
    read_bytes: Callable[..., Awaitable[bytes]],
    write_bytes: Callable[..., Awaitable[None]],
    accessor: Accessor | None = None,
    stdin: AsyncIterator[bytes] | bytes | None = None,
    prefix: str | PathSpec = "xx",
    digits: int = 2,
    suffix_format: str | None = None,
    keep_on_error: bool = False,
    silent: bool = False,
) -> tuple[ByteSource | None, IOResult]:
    if isinstance(prefix, PathSpec):
        prefix = prefix.mount_path
    suffix_fmt = suffix_format if suffix_format else f"%0{digits}d"
    if paths:
        raw = await read_bytes(accessor, paths[0])
    else:
        raw = await _read_stdin_async(stdin)
        if raw is None:
            raw = b""
    text = raw.decode(errors="replace")
    lines = split_lines(text)
    parts = _split_by_patterns(lines, list(patterns))
    writes: dict[str, bytes] = {}
    sizes: list[str] = []
    try:
        for idx, part in enumerate(parts):
            filename = prefix + (suffix_fmt % idx)
            data = ("\n".join(part) + "\n").encode() if part else b""
            await write_bytes(accessor, filename, data)
            writes[filename] = data
            sizes.append(str(len(data)))
    except Exception:
        if not keep_on_error:
            raise
    output = "" if silent or not sizes else "\n".join(sizes) + "\n"
    return output.encode(), IOResult(writes=writes)


__all__ = ["csplit"]