Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/streams/progress.py: 32%
66 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
1import logging
2from dataclasses import dataclass, field
4from anyio.abc import ObjectStream
5from tqdm import tqdm
7TQDM_KWARGS = {
8 "unit": "B",
9 "unit_scale": True,
10 "unit_divisor": 1024,
11}
13LOGGER = logging.getLogger(__name__)
16# Copied from https://github.com/tqdm/tqdm/pull/1172
17class logging_tqdm(tqdm):
18 """
19 A version of tqdm that outputs the progress bar
20 to Python logging instead of the console.
21 The progress will be logged with the info level.
22 Parameters
23 ----------
24 logger : logging.Logger, optional
25 Which logger to output to (default: logger.getLogger('tqdm.contrib.logging')).
26 All other parameters are passed on to regular tqdm,
27 with the following changed default:
28 mininterval: 1
29 bar_format: '{desc}{percentage:3.0f}%{r_bar}'
30 desc: 'progress: '
31 Example
32 -------
33 ```python
34 import logging
35 from time import sleep
36 from tqdm.contrib.logging import logging_tqdm
37 LOG = logging.getLogger(__name__)
38 if __name__ == '__main__':
39 logging.basicConfig(level=logging.INFO)
40 for _ in logging_tqdm(range(10), mininterval=1, logger=LOG):
41 sleep(0.3) # assume processing one item takes less than mininterval
42 ```
43 """
45 def __init__(
46 self,
47 *args,
48 # logger=None, # type: logging.Logger
49 # mininterval=1, # type: float
50 # bar_format='{desc}{percentage:3.0f}%{r_bar}', # type: str
51 # desc='progress: ', # type: str
52 **kwargs,
53 ):
54 if len(args) >= 2:
55 # Note: Due to Python 2 compatibility, we can't declare additional
56 # keyword arguments in the signature.
57 # As a result, we could get (due to the defaults below):
58 # TypeError: __init__() got multiple values for argument 'desc'
59 # This will raise a more descriptive error message.
60 # Calling dummy init to avoid attribute errors when __del__ is called
61 super(logging_tqdm, self).__init__([], disable=True)
62 raise ValueError("only iterable may be used as a positional argument")
63 tqdm_kwargs = kwargs.copy()
64 self._logger = tqdm_kwargs.pop("logger", None)
65 tqdm_kwargs.setdefault("mininterval", 1)
66 tqdm_kwargs.setdefault("bar_format", "{desc}{percentage:3.0f}%{r_bar}")
67 tqdm_kwargs.setdefault("desc", "progress: ")
68 self._last_log_n = -1
69 super(logging_tqdm, self).__init__(*args, **tqdm_kwargs)
71 def _get_logger(self):
72 if self._logger is not None:
73 return self._logger
74 return LOGGER
76 def display(self, msg=None, pos=None):
77 if not self.n:
78 # skip progress bar before having processed anything
79 LOGGER.debug("ignoring message before any progress: %r", self.n)
80 return
81 if self.n == self._last_log_n:
82 # avoid logging for the same progress multiple times
83 LOGGER.debug("ignoring log message with same n: %r", self.n)
84 return
85 self._last_log_n = self.n
86 if msg is None:
87 msg = self.__str__()
88 if not msg:
89 LOGGER.debug("ignoring empty message: %r", msg)
90 return
91 self._get_logger().info("%s", msg)
94@dataclass(kw_only=True)
95class ProgressStream(ObjectStream[bytes]):
96 stream: ObjectStream
97 logging: bool = False
99 __recv: tqdm = field(init=False, default=None)
100 __send: tqdm = field(init=False, default=None)
102 def __del__(self):
103 if self.__recv is not None:
104 self.__recv.close()
105 if self.__send is not None:
106 self.__send.close()
108 async def receive(self):
109 item = await self.stream.receive()
111 if self.__recv is None:
112 if self.logging:
113 self.__recv = logging_tqdm(desc="transfer", **TQDM_KWARGS)
114 else:
115 self.__recv = tqdm(desc="transfer", **TQDM_KWARGS)
117 self.__recv.update(len(item))
119 return item
121 async def send(self, item):
122 if self.__send is None:
123 if self.logging:
124 self.__send = logging_tqdm(desc="transfer", **TQDM_KWARGS)
125 else:
126 self.__send = tqdm(desc="transfer", **TQDM_KWARGS)
128 self.__send.update(len(item))
130 await self.stream.send(item)
132 async def send_eof(self):
133 await self.stream.send_eof()
135 async def aclose(self):
136 await self.stream.aclose()