"""Implements the generic progress logger class, and the ProgressBar class. """ from tqdm import tqdm, tqdm_notebook from collections import OrderedDict import time SETTINGS = { 'notebook': False } def notebook(turn='on'): SETTINGS['notebook'] = True if (turn == 'on') else False def troncate_string(s, max_length=25): return s if (len(s) < max_length) else (s[:max_length] + "...") class ProgressLogger: """Generic class for progress loggers. A progress logger contains a "state" dictionnary. Parameters ---------- init_state Dictionnary representing the initial state. """ def __init__(self, init_state=None): self.state = {} self.stored = {} self.logs = [] self.log_indent = 0 if init_state is not None: self.state.update(init_state) def log(self, message): self.logs.append((' ' * self.log_indent) + message) def dump_logs(self, filepath=None): if filepath is not None: with open(filepath, 'a') as f: f.write("\n".join(self.logs)) else: return "\n".join(self.logs) def callback(self, **kw): """Execute something after the state has been updated by the given state elements. This default callback does nothing, overwrite it by subclassing """ pass def store(self, **kw): """Store objects in the logger and trigger ``self.store_callback``. This works exactly like ``logger()``, but the later is meant for simple data objects (text, numbers) that will be sent over the network or written to a file. The ``store`` method expects rather large objects which are not necessarily serializable, and will be used eg to draw plots on the fly. """ self.stored.update(kw) self.store_callback(**kw) def store_callback(self, **kw): """Execute something after the store has been updated by the given state elements. This default callback does nothing, overwrite it by subclassing """ pass def iter(self, **kw): """Iterate through a list while updating the state. Examples -------- >>> for username in logger.iter(user=['tom', 'tim', 'lea']: >>> # At every loop, logger.state['user'] is updated >>> print (username) """ for field, iterable in kw.items(): for it in iterable: self(**{field: it}) yield it def __call__(self, **kw): self.state.update(kw) self.callback(**kw) class ProgressBarLogger(ProgressLogger): """Generic class for progress loggers. A progress logger contains a "state" dictionnary Parameters ---------- init_state Initial state of the logger bars Either None (will be initialized with no bar) or a list/tuple of bar names (``['main', 'sub']``) which will be initialized with index -1 and no total, or a dictionary (possibly ordered) of bars, of the form ``{bar_1: {title: 'bar1', index: 2, total:23}, bar_2: {...}}`` ignored_bars Either None (newly met bars will be added) or a list of blacklisted bar names, or ``'all_others'`` to signify that all bar names not already in ``self.bars`` will be ignored. """ bar_indent = 2 def __init__(self, init_state=None, bars=None, ignored_bars=None, logged_bars='all', min_time_interval=0, ignore_bars_under=0): ProgressLogger.__init__(self, init_state) if bars is None: bars = OrderedDict() elif isinstance(bars, (list, tuple)): bars = OrderedDict([ (b, dict(title=b, index=-1, total=None, message=None, indent=0)) for b in bars ]) if isinstance(ignored_bars, (list, tuple)): ignored_bars = set(ignored_bars) self.ignored_bars = ignored_bars self.logged_bars = logged_bars self.state['bars'] = bars self.min_time_interval = min_time_interval self.ignore_bars_under = ignore_bars_under @property def bars(self): """Return ``self.state['bars'].``""" return self.state['bars'] def bar_is_ignored(self, bar): if self.ignored_bars is None: return False elif self.ignored_bars == 'all_others': return (bar not in self.bars) else: return bar in self.ignored_bars def bar_is_logged(self, bar): if (not self.logged_bars): return False elif self.logged_bars == 'all': return True else: return bar in self.logged_bars def iterable_is_too_short(self, iterable): length = len(iterable) if hasattr(iterable, '__len__') else None return (length is not None) and (length < self.ignore_bars_under) def iter_bar(self, bar_prefix='', **kw): """Iterate through a list while updating a state bar. Examples -------- >>> for username in logger.iter_bar(user=['tom', 'tim', 'lea']): >>> # At every loop, logger.state['bars']['user'] is updated >>> # to {index: i, total: 3, title:'user'} >>> print (username) """ if 'bar_message' in kw: bar_message = kw.pop('bar_message') else: bar_message = None bar, iterable = kw.popitem() if self.bar_is_ignored(bar) or self.iterable_is_too_short(iterable): return iterable bar = bar_prefix + bar if hasattr(iterable, '__len__'): self(**{bar + '__total': len(iterable)}) def new_iterable(): last_time = time.time() i = 0 # necessary in case the iterator is empty for i, it in enumerate(iterable): now_time = time.time() if (i == 0) or (now_time - last_time > self.min_time_interval): if bar_message is not None: self(**{bar + '__message': bar_message(it)}) self(**{bar + '__index': i}) last_time = now_time yield it if self.bars[bar]['index'] != i: self(**{bar + '__index': i}) self(**{bar + '__index': i + 1}) return new_iterable() def bars_callback(self, bar, attr, value, old_value=None): """Execute a custom action after the progress bars are updated. Parameters ---------- bar Name/ID of the bar to be modified. attr Attribute of the bar attribute to be modified value New value of the attribute old_value Previous value of this bar's attribute. This default callback does nothing, overwrite it by subclassing. """ pass def __call__(self, **kw): items = sorted(kw.items(), key=lambda kv: not kv[0].endswith('total')) for key, value in items: if '__' in key: bar, attr = key.split('__') if self.bar_is_ignored(bar): continue kw.pop(key) if bar not in self.bars: self.bars[bar] = dict(title=bar, index=-1, total=None, message=None) old_value = self.bars[bar][attr] if self.bar_is_logged(bar): new_bar = (attr == 'index') and (value < old_value) if (attr == 'total') or (new_bar): self.bars[bar]['indent'] = self.log_indent else: self.log_indent = self.bars[bar]['indent'] self.log("[%s] %s: %s" % (bar, attr, value)) self.log_indent += self.bar_indent self.bars[bar][attr] = value self.bars_callback(bar, attr, value, old_value) self.state.update(kw) self.callback(**kw) class TqdmProgressBarLogger(ProgressBarLogger): """Tqdm-powered progress bar for console or Notebooks. Parameters ---------- init_state Initial state of the logger bars Either None (will be initialized with no bar) or a list/tuple of bar names (``['main', 'sub']``) which will be initialized with index -1 and no total, or a dictionary (possibly ordered) of bars, of the form ``{bar_1: {title: 'bar1', index: 2, total:23}, bar_2: {...}}`` ignored_bars Either None (newly met bars will be added) or a list of blacklisted bar names, or ``'all_others'`` to signify that all bar names not already in ``self.bars`` will be ignored. leave_bars notebook True will make the bars look nice (HTML) in the jupyter notebook. It is advised to leave to 'default' as the default can be globally set from inside a notebook with ``import proglog; proglog.notebook_mode()``. print_messages If True, every ``logger(message='something')`` will print a message in the console / notebook """ def __init__(self, init_state=None, bars=None, leave_bars=False, ignored_bars=None, logged_bars='all', notebook='default', print_messages=True, min_time_interval=0, ignore_bars_under=0): ProgressBarLogger.__init__(self, init_state=init_state, bars=bars, ignored_bars=ignored_bars, logged_bars=logged_bars, ignore_bars_under=ignore_bars_under, min_time_interval=min_time_interval) self.leave_bars = leave_bars self.tqdm_bars = OrderedDict([ (bar, None) for bar in self.bars ]) if notebook == 'default': notebook = SETTINGS['notebook'] self.notebook = notebook self.print_messages = print_messages self.tqdm = (tqdm_notebook if self.notebook else tqdm) def new_tqdm_bar(self, bar): """Create a new tqdm bar, possibly replacing an existing one.""" if (bar in self.tqdm_bars) and (self.tqdm_bars[bar] is not None): self.close_tqdm_bar(bar) infos = self.bars[bar] self.tqdm_bars[bar] = self.tqdm( total=infos['total'], desc=infos['title'], postfix=dict(now=troncate_string(str(infos['message']))), leave=self.leave_bars ) def close_tqdm_bar(self, bar): """Close and erase the tqdm bar""" self.tqdm_bars[bar].close() if not self.notebook: self.tqdm_bars[bar] = None def bars_callback(self, bar, attr, value, old_value): if (bar not in self.tqdm_bars) or (self.tqdm_bars[bar] is None): self.new_tqdm_bar(bar) if attr == 'index': if value >= old_value: total = self.bars[bar]['total'] if total and (value >= total): self.close_tqdm_bar(bar) else: self.tqdm_bars[bar].update(value - old_value) else: self.new_tqdm_bar(bar) self.tqdm_bars[bar].update(value + 1) elif attr == 'message': self.tqdm_bars[bar].set_postfix(now=troncate_string(str(value))) self.tqdm_bars[bar].update(0) def callback(self, **kw): if self.print_messages and ('message' in kw) and kw['message']: if self.notebook: print(kw['message']) else: self.tqdm.write(kw['message']) class RqWorkerProgressLogger: def __init__(self, job): self.job = job if 'progress_data' not in self.job.meta: self.job.meta['progress_data'] = {} self.job.save() def callback(self, **kw): self.job.meta['progress_data'] = self.state self.job.save() class RqWorkerBarLogger(RqWorkerProgressLogger, ProgressBarLogger): def __init__(self, job, init_state=None, bars=None, ignored_bars=(), logged_bars='all', min_time_interval=0): RqWorkerProgressLogger.__init__(self, job) ProgressBarLogger.__init__(self, init_state=init_state, bars=bars, ignored_bars=ignored_bars, logged_bars=logged_bars, min_time_interval=min_time_interval) class MuteProgressBarLogger(ProgressBarLogger): def bar_is_ignored(self, bar): return True def default_bar_logger(logger, bars=None, ignored_bars=None, logged_bars='all', min_time_interval=0, ignore_bars_under=0): if logger == 'bar': return TqdmProgressBarLogger( bars=bars, ignored_bars=ignored_bars, logged_bars=logged_bars, min_time_interval=min_time_interval, ignore_bars_under=ignore_bars_under ) elif logger is None: return MuteProgressBarLogger() else: return logger