aboutsummaryrefslogtreecommitdiff
path: root/.venv/lib/python3.12/site-packages/openai/cli/_progress.py
blob: 8a7f2525deb80aa9f24b22537cac725cefad4f4d (about) (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
from __future__ import annotations

import io
from typing import Callable
from typing_extensions import override


class CancelledError(Exception):
    def __init__(self, msg: str) -> None:
        self.msg = msg
        super().__init__(msg)

    @override
    def __str__(self) -> str:
        return self.msg

    __repr__ = __str__


class BufferReader(io.BytesIO):
    def __init__(self, buf: bytes = b"", desc: str | None = None) -> None:
        super().__init__(buf)
        self._len = len(buf)
        self._progress = 0
        self._callback = progress(len(buf), desc=desc)

    def __len__(self) -> int:
        return self._len

    @override
    def read(self, n: int | None = -1) -> bytes:
        chunk = io.BytesIO.read(self, n)
        self._progress += len(chunk)

        try:
            self._callback(self._progress)
        except Exception as e:  # catches exception from the callback
            raise CancelledError("The upload was cancelled: {}".format(e)) from e

        return chunk


def progress(total: float, desc: str | None) -> Callable[[float], None]:
    import tqdm

    meter = tqdm.tqdm(total=total, unit_scale=True, desc=desc)

    def incr(progress: float) -> None:
        meter.n = progress
        if progress == total:
            meter.close()
        else:
            meter.refresh()

    return incr


def MB(i: int) -> int:
    return int(i // 1024**2)