anystore.io

class IOFormat(StrEnum):
    """For use in typer cli"""

    csv = "csv"
    json = "json"

class ModelWriter(Writer):
    """
    A generic writer for pydantic objects to any out uri, either json or csv
    """

    def write(self, row: BaseModel) -> None:
        data = row.model_dump(by_alias=True, mode="json")
        return super().write(data)

class ProgressTask:
    """
    A labelled bar on a [`SyncProgressBar`][anystore.io.progress.SyncProgressBar].

    Use it as a context manager to have it disappear when its work is done:

    Example:
        ```python
        with bar.task("chunk 1", total=len(keys)) as task:
            for key in keys:
                task.advance(size=len(store.get(key)))
        ```
    """

    def __init__(
        self, bar: "SyncProgressBar", task_id: TaskID, throughput: Throughput
    ) -> None:
        self.bar = bar
        self.task_id = task_id
        self.throughput = throughput

    def advance(self, items: int | None = 1, size: int | None = None) -> None:
        """
        Advance the task, optionally accounting transferred bytes.

        Args:
            items: Number of items completed (default: 1)
            size: Bytes transferred for them, added to this task's throughput
                and to the overall throughput of the bar
        """
        if size:
            self.throughput.add(size)
            # a task can be fed the bar's own throughput to sum up the others
            if self.throughput is not self.bar.throughput:
                self.bar.throughput.add(size)
        self.bar.progress.advance(self.task_id, items or 0)

    def update(self, **kwargs: Any) -> None:
        """Update the underlying rich task, e.g. `total` or `description`"""
        self.bar.progress.update(self.task_id, **kwargs)

    def remove(self) -> None:
        """Remove the bar from the display"""
        self.bar.progress.remove_task(self.task_id)

    def __enter__(self) -> Self:
        return self

    def __exit__(self, *args: Any) -> None:
        self.remove()

class SmartHandler:
    def __init__(
        self,
        uri: Uri,
        compression: CompressKind | str | None = None,
        **kwargs: Any,
    ) -> None:
        self.uri = uri
        self.is_buffer = self.uri == "-"
        kwargs["mode"] = kwargs.get("mode", DEFAULT_MODE)
        self.mode = kwargs["mode"]
        self.compression = compression
        # stdio and an injected handle need the codec applied here; a uri goes
        # through `Store.open`, which applies it at the funnel – doing both
        # would encode the stream twice
        self.sys_io = _get_sysio(binary_mode(self.mode) if compression else self.mode)
        self.kwargs = kwargs
        self.handler: IO | None = None

    def open(self) -> IO[Any]:
        try:
            if self.is_buffer:
                return self._wrap(self.sys_io)
            elif isinstance(self.uri, (BytesIO, StringIO, IOBase)):
                return self._wrap(self.uri)
            else:
                resource = UriResource(self.uri)
                mode = self.kwargs.pop("mode", DEFAULT_MODE)
                self.handler = resource.open(
                    mode, compression=self.compression, **self.kwargs
                ).__enter__()
                return self.handler
        except FileNotFoundError as e:
            raise DoesNotExist(str(e))

    def _wrap(self, io: IO[Any]) -> IO[Any]:
        """Apply the codec to a handle this class did not open.

        `own=False`, so closing the codec flushes its frame and leaves the
        handle for whoever owns it – stdout, or the caller who passed it in.
        The codec itself does become ours to close, so it is tracked.
        """
        if self.compression is None:
            return io
        self.handler = open_codec(io, self.compression, self.mode)
        return self.handler

    def close(self):
        # a tracked handler is always ours to close: either the store's stream
        # or a codec we layered over someone else's handle
        if self.handler is not None:
            self.handler.close()

    def __enter__(self):
        return self.open()

    def __exit__(self, *args, **kwargs) -> None:
        self.close()

class SyncProgressBar:
    """
    A live display of one or more bars, with byte throughput.

    For a single bar, describe it up front and advance the display itself:

    Example:
        ```python
        with SyncProgressBar("download", total=len(keys)) as bar:
            for key in keys:
                bar.advance(size=len(store.get(key)))
        ```

    For concurrent work, add a bar per unit of it with
    [`task`][anystore.io.progress.SyncProgressBar.task] – they may be advanced
    from any thread, and all of them live in this one display, as wrapping each
    bar in a display of its own would leave all but the first invisible (rich
    only renders the outermost live display):

    Example:
        ```python
        with SyncProgressBar() as bar:
            for prefix in prefixes:
                with bar.task(prefix, total=len(keys)) as task:
                    for key in keys:
                        task.advance(size=len(store.get(key)))
        ```

    Args:
        description: Label for the display's own bar – leave it out for the
            multi-task case, where each task brings its own
        total: Number of items for that bar, if known
        console: Console to render on (default: a new one on the current
            `sys.stderr`)
        columns: Custom rich progress columns
        transient: Remove the whole display when it's done (default: yes)
        route_logging: Print log output through the same console while the
            display is running (default: yes)
        disable: Render nothing at all, e.g. for a `--quiet` flag
        window: Seconds to average throughput rates over
    """

    def __init__(
        self,
        description: str | None = None,
        total: int | None = None,
        console: Console | None = None,
        columns: Sequence[str | ProgressColumn] | None = None,
        transient: bool | None = True,
        route_logging: bool | None = True,
        disable: bool | None = False,
        window: int | None = DEFAULT_WINDOW,
    ) -> None:
        self.throughput = Throughput(window)
        self.description = description
        self.total = total
        self.console = console
        self.columns = columns or DEFAULT_COLUMNS
        self.transient = bool(transient)
        self.route_logging = bool(route_logging)
        self.disable = bool(disable)
        self.window = window
        self._progress: Progress | None = None
        self._task: ProgressTask | None = None
        self._stack = ExitStack()

    @property
    def progress(self) -> Progress:
        """The underlying rich `Progress`, once the display is started"""
        if self._progress is None:
            raise RuntimeError("Progress display not started")
        return self._progress

    @property
    def default_task(self) -> ProgressTask:
        """The display's own bar, for the single-bar case"""
        if self._task is None:
            raise RuntimeError(
                "No bar of its own, give the `SyncProgressBar` a description "
                "or add a `task()`"
            )
        return self._task

    def advance(self, items: int | None = 1, size: int | None = None) -> None:
        """
        Advance the display's own bar, optionally accounting transferred bytes.

        Only for the single-bar case – with several tasks, advance those.

        Args:
            items: Number of items completed (default: 1)
            size: Bytes transferred for them
        """
        self.default_task.advance(items, size)

    def update(self, **kwargs: Any) -> None:
        """Update the display's own bar, e.g. `total` or `description`"""
        self.default_task.update(**kwargs)

    def task(
        self,
        description: str,
        total: int | None = None,
        throughput: Throughput | None = None,
    ) -> ProgressTask:
        """
        Add a labelled bar to the display.

        Args:
            description: The label, e.g. the key prefix being worked on
            total: Number of items, if known – an unknown total pulses instead
                of filling up
            throughput: Byte counter to display, pass the bar's own
                `throughput` for a task that sums up all the others (default: a
                fresh one for this task)

        Returns:
            The task handle, removable and usable as a context manager
        """
        throughput = throughput if throughput is not None else Throughput(self.window)
        task_id = self.progress.add_task(
            description, total=total, throughput=throughput
        )
        return ProgressTask(self, task_id, throughput)

    def start(self) -> Self:
        """Start rendering (and, unless turned off, routing log output)"""
        # bind the console late: it holds on to the current `sys.stderr`, which
        # `logging_through` is about to replace
        console = self.console or Console(file=sys.stderr)
        self.console = console
        self._progress = Progress(
            *self.columns,
            console=console,
            transient=self.transient,
            disable=self.disable,
            # `logging_through` routes stderr through this console already, and
            # rich's own redirect would fight it
            redirect_stdout=False,
            redirect_stderr=False,
        )
        if self.route_logging:
            self._stack.enter_context(logging_through(console))
        self._stack.enter_context(self._progress)
        if self.description is not None or self.total is not None:
            # one counter for a single bar: the task shares the bar's own
            self._task = self.task(
                self.description or "", total=self.total, throughput=self.throughput
            )
        return self

    def stop(self) -> None:
        """Stop rendering and restore log output"""
        self._progress = None
        self._task = None
        self._stack.close()

    def __enter__(self) -> Self:
        return self.start()

    def __exit__(self, *args: Any) -> None:
        self.stop()

class Throughput:
    """
    Thread-safe byte counter exposing a rolling transfer rate.

    Args:
        window: Seconds to average the rate over
    """

    def __init__(self, window: int | None = DEFAULT_WINDOW) -> None:
        self.window = window or DEFAULT_WINDOW
        self.total = 0
        self._events: deque[tuple[float, int]] = deque()
        self._started = time.monotonic()
        self._lock = threading.Lock()

    def add(self, size: int) -> None:
        """Account `size` transferred bytes."""
        with self._lock:
            now = time.monotonic()
            self.total += size
            self._events.append((now, size))
            self._expire(now)

    @property
    def rate(self) -> float:
        """Bytes per second over the last `window` seconds."""
        with self._lock:
            now = time.monotonic()
            self._expire(now)
            span = min(now - self._started, self.window)
            if span <= 0:
                return 0.0
            return sum(size for _, size in self._events) / span

    def _expire(self, now: float) -> None:
        while self._events and now - self._events[0][0] > self.window:
            self._events.popleft()

class ThroughputColumn(ProgressColumn):
    """Renders the byte throughput of a task's `Throughput` field.

    The item count says how far along a task is; this says how fast bytes are
    actually moving.
    """

    def render(self, task: Task) -> Text:
        throughput: Throughput | None = task.fields.get("throughput")
        rate = throughput.rate if throughput is not None else 0.0
        return Text(format_bytes(rate, "B/s"), style="progress.data.speed")

class Writer:
    """
    A generic writer for python dict objects to any out uri, either json or csv

    Args:
        uri: string or path-like key uri to write to, `"-"` for stdout, or an
            already open handle
        mode: open mode, default `wb` (forced to text for csv)
        output_format: csv or json (default: json)
        fieldnames: csv header, inferred from the first row when omitted
        clean: Apply [clean_dict][anystore.util.data.clean_dict]
        compression: Codec to compress the output with ("gz", "bz2", "xz",
            "zst", "lz4")
        lazy: Defer creating the target to the first `write`, so a run that
            writes nothing leaves no file behind. Off by default – an empty
            csv carrying just its header is a legitimate thing to want.
        **kwargs: pass through storage-specific options
    """

    def __init__(
        self,
        uri: Uri,
        mode: str | None = DEFAULT_WRITE_MODE,
        output_format: Formats | None = "json",
        fieldnames: list[str] | None = None,
        clean: bool | None = False,
        compression: CompressKind | str | None = None,
        lazy: bool | None = False,
        **kwargs,
    ) -> None:
        if output_format not in (FORMAT_JSON, FORMAT_CSV):
            raise ValueError("Invalid output format, only csv or json allowed")
        mode = mode or DEFAULT_WRITE_MODE
        self.mode = mode.replace("b", "") if output_format == "csv" else mode
        self.handler = SmartHandler(
            uri, mode=self.mode, compression=compression, **kwargs
        )
        self.fieldnames = fieldnames
        self.output_format = output_format
        self.clean = clean
        self.lazy = lazy
        self.csv_writer: csv.DictWriter | None = None
        self._io: IO[Any] | None = None

    def open(self) -> IO[Any]:
        """The open target, created on first ask.

        A csv whose header is known up front writes it here, so a writer that
        never sees a row still leaves a complete – if empty – table instead of
        a zero-byte file.
        """
        if self._io is None:
            self._io = self.handler.open()
            if self.output_format == FORMAT_CSV and self.fieldnames:
                self._start_csv(self.fieldnames)
        return self._io

    def _start_csv(self, fieldnames: Iterable[str]) -> None:
        """Bind the csv writer to the open target and write its header."""
        self.csv_writer = csv.DictWriter(self.io, fieldnames)
        self.csv_writer.writeheader()

    def close(self) -> None:
        """Flush and close the target, if one was ever opened."""
        self.handler.close()

    @property
    def io(self) -> IO[Any]:
        return self.open()

    def __enter__(self) -> Self:
        if not self.lazy:
            self.open()
        return self

    def __exit__(self, *args) -> None:
        self.close()

    def write(self, row: SDict) -> None:
        if self.output_format == FORMAT_CSV:
            # opening starts the writer itself when the header is known up front
            self.open()
            if self.csv_writer is None:
                self._start_csv(row.keys())

        if self.output_format == FORMAT_JSON:
            if self.clean:
                row = clean_dict(row)
            line = orjson.dumps(
                row,
                default=_default_serializer,
                option=orjson.OPT_APPEND_NEWLINE | orjson.OPT_NAIVE_UTC,
            )
            if "b" not in self.mode:
                line = line.decode()
            self.io.write(line)
        elif self.csv_writer:
            self.csv_writer.writerow(row)