from contextlib import contextmanager
import os
import shutil
import sys
import time
import logging
import inspect
import pprint
import subprocess
import textwrap
[docs]class Echo:
""" Context maganger for echoing variable assignments (in CPython) """
def __init__(self, msg, indent=' '):
self.msg = msg
self.indent = indent
self.parent_frame = inspect.currentframe().f_back
def __enter__(self):
print(self.msg)
self.locals_on_entry = self.parent_frame.f_locals.copy()
def __exit__(self, exc_t, exc_v, tb):
new_locals = dict((k, v) for k, v in self.parent_frame.f_locals.items() if k not in self.locals_on_entry)
print(textwrap.indent(pprint.pformat(new_locals), self.indent))
[docs]class Notify:
def __init__(self, msg=lambda t: "Job finished in %.1f seconds" % t):
self.msg = msg
self.t0 = time.time()
def __enter__(self):
pass
[docs] def notify(self, title, message):
if os.environ.get('DISPLAY', ''):
subprocess.call(['notify-send', title, message])
else:
print(title)
print(message)
def __exit__(self, exc_t, exc_v, tb):
if exc_t is None:
title = "Success"
else:
title = "Failure"
self.notify(title, self.msg(time.time() - self.t0))
[docs]class c:
ok = '\033[92m' # green
fail = '\033[91m' # red
endc = '\033[0m' # reset
[docs]class Timed:
""" Utility function for timing portions of your python script
Parameters
----------
msg : str
timer : callable
You can switch to e.g. ``time.process_time`` but note that if other
programs are called from e.g. ``subprocess.Popen`` the time spent
in those subprocesses will not be included.
Examples
--------
>>> t = Timed("Counting stars...").tic(); stars.count(); t.toc_and_print() # doctest: +SKIP
Counting stars... (42.0 s) [ ok]
>>> with Timed("Counting sheep..."): # doctest: +SKIP
... n_sheep = animals.count('sheep')
...
Counting sheep... (17.2 s) [ ok]
"""
counting = False
def __init__(self, msg=None, timer=time.time, fmt_s='.1f', out=sys.stdout):
self.msg = msg
self.out = out
self.fmt_s = fmt_s
sys.stdout.flush()
self.timer = timer
[docs] def tic(self):
if self.msg is not None:
self.out.write(self.msg)
self.out.flush()
self.counting = True
self.t = self.timer()
return self # Allows t = Timed().tic(); integrals.calc(); t = t.toc()
[docs] def toc(self, ok=True):
if self.counting:
t = self.timer() - self.t
if self.msg is not None:
if ok:
status = 'ok'
color = c.ok
else:
status = 'error'
color = c.fail
self.out.write('%{}s\n'.format(shutil.get_terminal_size()[0] - len(self.msg)) %
('(%{fmt_s} s) [{c}%5s{r}]'.format(fmt_s=self.fmt_s, c=color, r=c.endc) % (t, status)))
self.out.flush()
return t
else:
raise ValueError("Not counting, did you forget to call ``.tic()`` method?")
def __enter__(self):
self.tic()
def __exit__(self, exc_type, exc_value, traceback):
self.toc(not exc_type)
[docs]@contextmanager
def limit_logging(max_lvl=logging.CRITICAL):
""" Contextmanager for silencing logging messages.
Examples
--------
>>> with limit_logging():
... logger.info("you won't see this...") # doctest: +SKIP
"""
_ori = logging.root.manager.disable
logging.disable(max_lvl)
try:
yield
finally:
logging.disable(_ori)