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

1import logging 

2from dataclasses import dataclass, field 

3 

4from anyio.abc import ObjectStream 

5from tqdm import tqdm 

6 

7TQDM_KWARGS = { 

8 "unit": "B", 

9 "unit_scale": True, 

10 "unit_divisor": 1024, 

11} 

12 

13LOGGER = logging.getLogger(__name__) 

14 

15 

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 """ 

44 

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) 

70 

71 def _get_logger(self): 

72 if self._logger is not None: 

73 return self._logger 

74 return LOGGER 

75 

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) 

92 

93 

94@dataclass(kw_only=True) 

95class ProgressStream(ObjectStream[bytes]): 

96 stream: ObjectStream 

97 logging: bool = False 

98 

99 __recv: tqdm = field(init=False, default=None) 

100 __send: tqdm = field(init=False, default=None) 

101 

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() 

107 

108 async def receive(self): 

109 item = await self.stream.receive() 

110 

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) 

116 

117 self.__recv.update(len(item)) 

118 

119 return item 

120 

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) 

127 

128 self.__send.update(len(item)) 

129 

130 await self.stream.send(item) 

131 

132 async def send_eof(self): 

133 await self.stream.send_eof() 

134 

135 async def aclose(self): 

136 await self.stream.aclose()