Skip to content

Commit

Permalink
std: use SIGWINCH signal to update the size
Browse files Browse the repository at this point in the history
SIGWINCH is a signal send to the process when the host terminal is
resized, note that it's only available on Unix platforms. Catching this
signal allows us to dynamically adapt the number of columns at a very
low cost compared to `dynamic_ncols`. Indeed, as this approach is signal
based, there is no need to poll the terminal size on each iteration.
  • Loading branch information
LucasChollet committed Dec 16, 2023
1 parent 4c956c2 commit b8e0dc9
Showing 1 changed file with 34 additions and 9 deletions.
43 changes: 34 additions & 9 deletions tqdm/std.py
Expand Up @@ -7,6 +7,7 @@
>>> for i in trange(10):
... ...
"""
import signal
import sys
from collections import OrderedDict, defaultdict
from contextlib import contextmanager
Expand Down Expand Up @@ -367,6 +368,8 @@ class tqdm(Comparable):
monitor = None
_instances = WeakSet()

registered_classes = set()

@staticmethod
def format_sizeof(num, suffix='', divisor=1000):
"""
Expand Down Expand Up @@ -665,6 +668,7 @@ def format_meter(n, total, elapsed, ncols=None, prefix='', ascii=False, unit='it

def __new__(cls, *_, **__):
instance = object.__new__(cls)
tqdm.registered_classes.add(cls)
with cls.get_lock(): # also constructs lock if non-existent
cls._instances.add(instance)
# create monitoring thread
Expand Down Expand Up @@ -1008,17 +1012,18 @@ def __init__(self, iterable=None, desc=None, total=None, leave=True, file=None,
TqdmKeyError("Unknown argument(s): " + str(kwargs)))

# Preprocess the arguments
keep_original_size = ncols is not None, nrows is not None
force_dynamic_ncols_update = dynamic_ncols
if (
(ncols is None or nrows is None) and (file in (sys.stderr, sys.stdout))
) or dynamic_ncols: # pragma: no cover
if dynamic_ncols:
dynamic_ncols = _screen_shape_wrapper()
if dynamic_ncols:
ncols, nrows = dynamic_ncols(file)
) or force_dynamic_ncols_update: # pragma: no cover
dynamic_ncols = _screen_shape_wrapper()
if force_dynamic_ncols_update and dynamic_ncols:
keep_original_size = False, False
ncols, nrows = dynamic_ncols(file)
else:
_dynamic_ncols = _screen_shape_wrapper()
if _dynamic_ncols:
_ncols, _nrows = _dynamic_ncols(file)
if dynamic_ncols:
_ncols, _nrows = dynamic_ncols(file)
if ncols is None:
ncols = _ncols
if nrows is None:
Expand Down Expand Up @@ -1054,6 +1059,7 @@ def __init__(self, iterable=None, desc=None, total=None, leave=True, file=None,
self.fp = file
self.ncols = ncols
self.nrows = nrows
self.keep_original_size = keep_original_size
self.mininterval = mininterval
self.maxinterval = maxinterval
self.miniters = miniters
Expand All @@ -1067,6 +1073,7 @@ def __init__(self, iterable=None, desc=None, total=None, leave=True, file=None,
self.lock_args = lock_args
self.delay = delay
self.gui = gui
self.force_dynamic_ncols_update = force_dynamic_ncols_update
self.dynamic_ncols = dynamic_ncols
self.smoothing = smoothing
self._ema_dn = EMA(smoothing)
Expand Down Expand Up @@ -1450,7 +1457,7 @@ def format_dict(self):
if self.disable and not hasattr(self, 'unit'):
return defaultdict(lambda: None, {
'n': self.n, 'total': self.total, 'elapsed': 0, 'unit': 'it'})
if self.dynamic_ncols:
if self.force_dynamic_ncols_update and self.dynamic_ncols:
self.ncols, self.nrows = self.dynamic_ncols(self.fp)
return {
'n': self.n, 'total': self.total,
Expand Down Expand Up @@ -1523,3 +1530,21 @@ def wrapattr(cls, stream, method, total=None, bytes=True, **tqdm_kwargs):
def trange(*args, **kwargs):
"""Shortcut for tqdm(range(*args), **kwargs)."""
return tqdm(range(*args), **kwargs)


def resize_signal_handler(signalnum, frame):
for cls in tqdm.registered_classes:
with cls.get_lock():
for instance in cls._instances:
if instance.dynamic_ncols:
ncols, nrows = instance.dynamic_ncols(instance.fp)
if not instance.keep_original_size[0]:
instance.ncols = ncols
if not instance.keep_original_size[1]:
instance.nrows = nrows


try:
signal.signal(signal.SIGWINCH, resize_signal_handler)
except AttributeError:
pass # Some systems, like Windows, do not have SIGWINCH

0 comments on commit b8e0dc9

Please sign in to comment.