Source code for glompo.common.wrappers

""" Decorators and wrappers used throughout GloMPO. """

import importlib
import inspect
import sys
import warnings
from functools import wraps
from pathlib import Path


[docs]def process_print_redirect(opt_id, working_dir, func): """ Redirects a process' output to a text file in a designated directory. """ @wraps(func) def wrapper(*args, **kwargs): sys.stdout = Path(working_dir, "glompo_optimizer_printstreams", f"printstream_{opt_id:04}.out").open("w") sys.stderr = Path(working_dir, "glompo_optimizer_printstreams", f"printstream_{opt_id:04}.err").open("w") func(*args, **kwargs) sys.stdout.close() sys.stderr.close() return wrapper
[docs]def catch_user_interrupt(func): """ Catches a :exc:`KeyboardInterrupt` signal and exits gracefully. """ @wraps(func) def wrapper(*args, **kwargs): try: return func(*args, **kwargs) except KeyboardInterrupt: print("Interrupt signal received. Process stopping.") return wrapper
[docs]def decorate_all_methods(decorator): """ Applies `decorator` to every method in a class. """ def apply_decorator(cls): for key, func in cls.__dict__.items(): if inspect.isfunction(func): setattr(cls, key, decorator(func)) return cls return apply_decorator
[docs]def needs_optional_package(package: str): """ Checks `package` requirement before running a function. Will warn and not allow a function to execute if the `package` it requires is not available to the system. Applied to functions which provide optional GloMPO functionality. """ def decorator(func): @wraps(func) def wrapper(*args, **kwargs): try: importlib.import_module(package) return func(*args, **kwargs) except (ModuleNotFoundError, ImportError): warnings.warn(f"Unable to construct checkpoint without {package} package installed.", ResourceWarning) return None return wrapper return decorator