Source code for k1lib._context

# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
import io, sys, warnings, time, k1lib, os
from contextlib import contextmanager
from functools import partial
try: import rdkit; hasRdkit = True
except: hasRdkit = False
try: import matplotlib.pyplot as plt
except: pass
__all__ = ["captureStdout", "capturePlt", "ignoreWarnings", "timer", "attrContext"]
[docs]@contextmanager def captureStdout(out=True, c=False) -> k1lib.Wrapper: # captureStdout """Captures every print() statement. Taken from https://stackoverflow.com/questions/16571150/how-to-capture-stdout-output-from-a-python-function-call. Example:: with k1lib.captureStdout() as outer: print("something") with k1lib.captureStdout() as inner: print("inside inner") print("else") # prints "['something', 'else']" print(outer.value) # prints "['inside inner']" print(inner.value) Note that internally, this replaces :data:`sys.stdout` as :class:`io.StringIO`, so might not work property if you have fancy :class:`bytes` stuff going on. Also, carriage return (``\\r``) will erase the line, so multi-line overlaps might not show up correctly. If you wish to still print stuff out, you can do something like this:: with k1.captureStdout() as out: print("abc") # gets captured out.print("def") # not captured, will actually print out # you can also access the stream like this print("something", file=out.old) Also, by default, this won't work if you're trying to capture C library's output, because they write to stdout directly, instead of going through Python's mechanism. You can capture it down at the C level by doing this:: with captureStdout() as out1, captureStdout(c=True) as out2: os.system("ls") # gets captured by out2, but not out1 print("abc") # gets captured by out1, but not out2 It's a little bit hard to actually integrate C mode and non-C mode together, so for the time being, you gotta have 2 context managers if you want to capture absolutely everything, C or not. :param out: if True, captures stdout, else captures stderr :param c: whether to capture at the C/C++ level or not""" # captureStdout if c is True: # captureStdout import wurlitzer; w = k1lib.Wrapper("") # captureStdout try: # captureStdout with wurlitzer.pipes() as (_out, _err): yield w # captureStdout except BrokenPipeError: pass # captureStdout finally: w.value = _out.read().split("\n") if out else _err.read().split("\n") # captureStdout else: # captureStdout if out: _stdout = sys.stdout; sys.stdout = _stringio = io.StringIO() # captureStdout else: _stdout = sys.stderr; sys.stderr = _stringio = io.StringIO() # captureStdout w = k1lib.Wrapper([]) # captureStdout w.print = partial(print, file=_stdout) # captureStdout w.old = _stdout # captureStdout try: yield w # captureStdout finally: # captureStdout w.value = [l.split("\r")[-1] for l in _stringio.getvalue().split("\n")] # captureStdout if out: sys.stdout = _stdout # captureStdout else: sys.stderr = _stdout # captureStdout
[docs]@contextmanager # captureStdout def capturePlt(): # capturePlt """Tries to capture matplotlib plot. Example:: x = np.linspace(-2, 2) with k1.capturePlt() as fig: plt.plot(x, x**2) plt.show() capturedImage = fig() # reference it here This is a convenience function to deal with libraries that call ``plt.show()`` and doesn't let us intercept at the middle to generate an image.""" # capturePlt try: # capturePlt ans = k1lib.Wrapper(None) # capturePlt with k1lib._settings.monkey.context(capturePlt=True): # capturePlt yield ans # capturePlt ans.value = plt._k1_capturedImg() # capturePlt finally: pass # capturePlt
[docs]@contextmanager # capturePlt def ignoreWarnings(): # ignoreWarnings """Context manager to ignore every warning. Example:: import warnings with k1lib.ignoreWarnings(): warnings.warn("some random stuff") # will not show anything""" # ignoreWarnings with warnings.catch_warnings(): # ignoreWarnings warnings.simplefilter("ignore") # ignoreWarnings warnings.filterwarnings('ignore') # ignoreWarnings if hasRdkit: rdkit.RDLogger.DisableLog('rdApp.*') # ignoreWarnings yield # ignoreWarnings if hasRdkit: rdkit.RDLogger.EnableLog('rdApp.*') # ignoreWarnings
[docs]@contextmanager # ignoreWarnings def timer(): # timer """Times generic code. Example:: with k1lib.timer() as t: time.sleep(1.1) # prints out float close to 1.1 print(t()) The with- statement will actually return a :class:`~k1lib.Wrapper` with value None. The correct time will be deposited into it after the code block ends.""" # timer w = k1lib.Wrapper(None) # timer beginTime = time.time() # timer try: yield w # timer finally: w.value = time.time() - beginTime # timer
[docs]@contextmanager # timer def attrContext(var, **kwargs): # attrContext """Temporarily sets variable's attribute to something else. Example:: class A: pass a = A() a.b = 3 print(a.b) # prints "3" with k1lib.attrContext(a, b=4, c=5): print(a.b, a.c) # prints "4 5" print(a.b, a.c) # prints "3 None" """ # attrContext oldValues = dict() # attrContext for k, v in kwargs.items(): # attrContext oldValues[k] = getattr(var, k, None) # attrContext setattr(var, k, v) # attrContext try: yield # attrContext finally: # attrContext for k, v in oldValues.items(): # attrContext setattr(var, k, v) # attrContext