Coverage for /Volumes/workspace/python-progressbar/.tox/pypy3/site-packages/progressbar/utils.py: 93%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import atexit
4import datetime
5import io
6import logging
7import os
8import re
9import sys
11from python_utils import types
12from python_utils.converters import scale_1024
13from python_utils.terminal import get_terminal_size
14from python_utils.time import epoch
15from python_utils.time import format_time
16from python_utils.time import timedelta_to_seconds
18if types.TYPE_CHECKING:
19 from .bar import ProgressBar
21assert timedelta_to_seconds
22assert get_terminal_size
23assert format_time
24assert scale_1024
25assert epoch
27ANSI_TERMS = (
28 '([xe]|bv)term',
29 '(sco)?ansi',
30 'cygwin',
31 'konsole',
32 'linux',
33 'rxvt',
34 'screen',
35 'tmux',
36 'vt(10[02]|220|320)',
37)
38ANSI_TERM_RE = re.compile('^({})'.format('|'.join(ANSI_TERMS)), re.IGNORECASE)
41def is_ansi_terminal(fd: types.IO, is_terminal: bool | None = None) \
42 -> bool: # pragma: no cover
43 if is_terminal is None:
44 # Jupyter Notebooks define this variable and support progress bars
45 if 'JPY_PARENT_PID' in os.environ:
46 is_terminal = True
47 # This works for newer versions of pycharm only. older versions there
48 # is no way to check.
49 elif os.environ.get('PYCHARM_HOSTED') == '1':
50 is_terminal = True
52 if is_terminal is None:
53 # check if we are writing to a terminal or not. typically a file object
54 # is going to return False if the instance has been overridden and
55 # isatty has not been defined we have no way of knowing so we will not
56 # use ansi. ansi terminals will typically define one of the 2
57 # environment variables.
58 try:
59 is_tty = fd.isatty()
60 # Try and match any of the huge amount of Linux/Unix ANSI consoles
61 if is_tty and ANSI_TERM_RE.match(os.environ.get('TERM', '')):
62 is_terminal = True
63 # ANSICON is a Windows ANSI compatible console
64 elif 'ANSICON' in os.environ:
65 is_terminal = True
66 else:
67 is_terminal = None
68 except Exception:
69 is_terminal = False
71 return is_terminal
74def is_terminal(fd: types.IO, is_terminal: bool | None = None) -> bool:
75 if is_terminal is None:
76 # Full ansi support encompasses what we expect from a terminal
77 is_terminal = is_ansi_terminal(True) or None
79 if is_terminal is None:
80 # Allow a environment variable override
81 is_terminal = env_flag('PROGRESSBAR_IS_TERMINAL', None)
83 if is_terminal is None: # pragma: no cover
84 # Bare except because a lot can go wrong on different systems. If we do
85 # get a TTY we know this is a valid terminal
86 try:
87 is_terminal = fd.isatty()
88 except Exception:
89 is_terminal = False
91 return is_terminal
94def deltas_to_seconds(*deltas,
95 **kwargs) -> int | float | None: # default=ValueError):
96 '''
97 Convert timedeltas and seconds as int to seconds as float while coalescing
99 >>> deltas_to_seconds(datetime.timedelta(seconds=1, milliseconds=234))
100 1.234
101 >>> deltas_to_seconds(123)
102 123.0
103 >>> deltas_to_seconds(1.234)
104 1.234
105 >>> deltas_to_seconds(None, 1.234)
106 1.234
107 >>> deltas_to_seconds(0, 1.234)
108 0.0
109 >>> deltas_to_seconds()
110 Traceback (most recent call last):
111 ...
112 ValueError: No valid deltas passed to `deltas_to_seconds`
113 >>> deltas_to_seconds(None)
114 Traceback (most recent call last):
115 ...
116 ValueError: No valid deltas passed to `deltas_to_seconds`
117 >>> deltas_to_seconds(default=0.0)
118 0.0
119 '''
120 default = kwargs.pop('default', ValueError)
121 assert not kwargs, 'Only the `default` keyword argument is supported'
123 for delta in deltas:
124 if delta is None:
125 continue
126 if isinstance(delta, datetime.timedelta):
127 return timedelta_to_seconds(delta)
128 elif not isinstance(delta, float):
129 return float(delta)
130 else:
131 return delta
133 if default is ValueError:
134 raise ValueError('No valid deltas passed to `deltas_to_seconds`')
135 else:
136 return default
139def no_color(value: types.StringTypes) -> types.StringTypes:
140 '''
141 Return the `value` without ANSI escape codes
143 >>> no_color(b'\u001b[1234]abc') == b'abc'
144 True
145 >>> str(no_color(u'\u001b[1234]abc'))
146 'abc'
147 >>> str(no_color('\u001b[1234]abc'))
148 'abc'
149 '''
150 if isinstance(value, bytes):
151 pattern = '\\\u001b\\[.*?[@-~]'
152 pattern = pattern.encode()
153 replace = b''
154 assert isinstance(pattern, bytes)
155 else:
156 pattern = u'\x1b\\[.*?[@-~]'
157 replace = ''
159 return re.sub(pattern, replace, value)
162def len_color(value: types.StringTypes) -> int:
163 '''
164 Return the length of `value` without ANSI escape codes
166 >>> len_color(b'\u001b[1234]abc')
167 3
168 >>> len_color(u'\u001b[1234]abc')
169 3
170 >>> len_color('\u001b[1234]abc')
171 3
172 '''
173 return len(no_color(value))
176def env_flag(name: str, default: bool | None = None) -> bool | None:
177 '''
178 Accepts environt variables formatted as y/n, yes/no, 1/0, true/false,
179 on/off, and returns it as a boolean
181 If the environment variable is not defined, or has an unknown value,
182 returns `default`
183 '''
184 v = os.getenv(name)
185 if v and v.lower() in ('y', 'yes', 't', 'true', 'on', '1'):
186 return True
187 if v and v.lower() in ('n', 'no', 'f', 'false', 'off', '0'):
188 return False
189 return default
192class WrappingIO:
194 def __init__(self, target: types.IO, capturing: bool = False,
195 listeners: types.Set[ProgressBar] = None) -> None:
196 self.buffer = io.StringIO()
197 self.target = target
198 self.capturing = capturing
199 self.listeners = listeners or set()
200 self.needs_clear = False
202 def isatty(self): # pragma: no cover
203 return self.target.isatty()
205 def write(self, value: str) -> None:
206 if self.capturing:
207 self.buffer.write(value)
208 if '\n' in value: # pragma: no branch
209 self.needs_clear = True
210 for listener in self.listeners: # pragma: no branch
211 listener.update()
212 else:
213 self.target.write(value)
214 if '\n' in value: # pragma: no branch
215 self.flush_target()
217 def flush(self) -> None:
218 self.buffer.flush()
220 def _flush(self) -> None:
221 value = self.buffer.getvalue()
222 if value:
223 self.flush()
224 self.target.write(value)
225 self.buffer.seek(0)
226 self.buffer.truncate(0)
227 self.needs_clear = False
229 # when explicitly flushing, always flush the target as well
230 self.flush_target()
232 def flush_target(self) -> None: # pragma: no cover
233 if not self.target.closed and getattr(self.target, 'flush'):
234 self.target.flush()
237class StreamWrapper:
238 '''Wrap stdout and stderr globally'''
240 def __init__(self):
241 self.stdout = self.original_stdout = sys.stdout
242 self.stderr = self.original_stderr = sys.stderr
243 self.original_excepthook = sys.excepthook
244 self.wrapped_stdout = 0
245 self.wrapped_stderr = 0
246 self.wrapped_excepthook = 0
247 self.capturing = 0
248 self.listeners = set()
250 if env_flag('WRAP_STDOUT', default=False): # pragma: no cover
251 self.wrap_stdout()
253 if env_flag('WRAP_STDERR', default=False): # pragma: no cover
254 self.wrap_stderr()
256 def start_capturing(self, bar: ProgressBar | None = None) -> None:
257 if bar: # pragma: no branch
258 self.listeners.add(bar)
260 self.capturing += 1
261 self.update_capturing()
263 def stop_capturing(self, bar: ProgressBar | None = None) -> None:
264 if bar: # pragma: no branch
265 try:
266 self.listeners.remove(bar)
267 except KeyError:
268 pass
270 self.capturing -= 1
271 self.update_capturing()
273 def update_capturing(self) -> None: # pragma: no cover
274 if isinstance(self.stdout, WrappingIO):
275 self.stdout.capturing = self.capturing > 0
277 if isinstance(self.stderr, WrappingIO):
278 self.stderr.capturing = self.capturing > 0
280 if self.capturing <= 0:
281 self.flush()
283 def wrap(self, stdout: bool = False, stderr: bool = False) -> None:
284 if stdout:
285 self.wrap_stdout()
287 if stderr:
288 self.wrap_stderr()
290 def wrap_stdout(self) -> types.IO:
291 self.wrap_excepthook()
293 if not self.wrapped_stdout:
294 self.stdout = sys.stdout = WrappingIO(self.original_stdout,
295 listeners=self.listeners)
296 self.wrapped_stdout += 1
298 return sys.stdout
300 def wrap_stderr(self) -> types.IO:
301 self.wrap_excepthook()
303 if not self.wrapped_stderr:
304 self.stderr = sys.stderr = WrappingIO(self.original_stderr,
305 listeners=self.listeners)
306 self.wrapped_stderr += 1
308 return sys.stderr
310 def unwrap_excepthook(self) -> None:
311 if self.wrapped_excepthook:
312 self.wrapped_excepthook -= 1
313 sys.excepthook = self.original_excepthook
315 def wrap_excepthook(self) -> None:
316 if not self.wrapped_excepthook:
317 logger.debug('wrapping excepthook')
318 self.wrapped_excepthook += 1
319 sys.excepthook = self.excepthook
321 def unwrap(self, stdout: bool = False, stderr: bool = False) -> None:
322 if stdout:
323 self.unwrap_stdout()
325 if stderr:
326 self.unwrap_stderr()
328 def unwrap_stdout(self) -> None:
329 if self.wrapped_stdout > 1:
330 self.wrapped_stdout -= 1
331 else:
332 sys.stdout = self.original_stdout
333 self.wrapped_stdout = 0
335 def unwrap_stderr(self) -> None:
336 if self.wrapped_stderr > 1:
337 self.wrapped_stderr -= 1
338 else:
339 sys.stderr = self.original_stderr
340 self.wrapped_stderr = 0
342 def needs_clear(self) -> bool: # pragma: no cover
343 stdout_needs_clear = getattr(self.stdout, 'needs_clear', False)
344 stderr_needs_clear = getattr(self.stderr, 'needs_clear', False)
345 return stderr_needs_clear or stdout_needs_clear
347 def flush(self) -> None:
348 if self.wrapped_stdout: # pragma: no branch
349 try:
350 self.stdout._flush()
351 except (io.UnsupportedOperation,
352 AttributeError): # pragma: no cover
353 self.wrapped_stdout = False
354 logger.warn('Disabling stdout redirection, %r is not seekable',
355 sys.stdout)
357 if self.wrapped_stderr: # pragma: no branch
358 try:
359 self.stderr._flush()
360 except (io.UnsupportedOperation,
361 AttributeError): # pragma: no cover
362 self.wrapped_stderr = False
363 logger.warn('Disabling stderr redirection, %r is not seekable',
364 sys.stderr)
366 def excepthook(self, exc_type, exc_value, exc_traceback):
367 self.original_excepthook(exc_type, exc_value, exc_traceback)
368 self.flush()
371class AttributeDict(dict):
372 '''
373 A dict that can be accessed with .attribute
375 >>> attrs = AttributeDict(spam=123)
377 # Reading
379 >>> attrs['spam']
380 123
381 >>> attrs.spam
382 123
384 # Read after update using attribute
386 >>> attrs.spam = 456
387 >>> attrs['spam']
388 456
389 >>> attrs.spam
390 456
392 # Read after update using dict access
394 >>> attrs['spam'] = 123
395 >>> attrs['spam']
396 123
397 >>> attrs.spam
398 123
400 # Read after update using dict access
402 >>> del attrs.spam
403 >>> attrs['spam']
404 Traceback (most recent call last):
405 ...
406 KeyError: 'spam'
407 >>> attrs.spam
408 Traceback (most recent call last):
409 ...
410 AttributeError: No such attribute: spam
411 >>> del attrs.spam
412 Traceback (most recent call last):
413 ...
414 AttributeError: No such attribute: spam
415 '''
417 def __getattr__(self, name: str) -> int:
418 if name in self:
419 return self[name]
420 else:
421 raise AttributeError("No such attribute: " + name)
423 def __setattr__(self, name: str, value: int) -> None:
424 self[name] = value
426 def __delattr__(self, name: str) -> None:
427 if name in self:
428 del self[name]
429 else:
430 raise AttributeError("No such attribute: " + name)
433logger = logging.getLogger(__name__)
434streams = StreamWrapper()
435atexit.register(streams.flush)