Source code for mdevaluate.autosave

import os
import numpy as np
import functools
import inspect

from .checksum import checksum
from .logging import logger

autosave_directory = None
load_autosave_data = False
verbose_print = True
user_autosave_directory = os.path.join(os.environ['HOME'], '.mdevaluate/autosave')


[docs]def notify(msg): if verbose_print: logger.info(msg) else: logger.debug(msg)
[docs]def enable(dir, load_data=True, verbose=True): """ Enable auto saving results of functions decorated with :func:`autosave_data`. Args: dir: Directory where the data should be saved. load_data (opt., bool): If data should also be loaded. """ global autosave_directory, load_autosave_data, verbose_print verbose_print = verbose # absolute = os.path.abspath(dir) # os.makedirs(absolute, exist_ok=True) autosave_directory = dir load_autosave_data = load_data notify('Enabled autosave in directory: {}'.format(autosave_directory))
[docs]def disable(): """Disable autosave.""" global autosave_directory, load_autosave_data autosave_directory = None load_autosave_data = False
[docs]class disabled: """ A context manager that disbales the autosave module within its context. Example: import mdevaluate as md md.autosave.enable('data') with md.autosave.disabled(): # Autosave functionality is disabled within this context. md.correlation.shifted_correlation( ... ) # After the context is exited, autosave will work as before. """ def __enter__(self): self._autosave_directory = autosave_directory disable() def __exit__(self, *args): enable(self._autosave_directory)
[docs]def get_directory(reader): """Get the autosave directory for a trajectory reader.""" outdir = os.path.dirname(reader.filename) savedir = os.path.join(outdir, autosave_directory) if not os.path.exists(savedir): try: os.makedirs(savedir) except PermissionError: pass if not os.access(savedir, os.W_OK): savedir = os.path.join(user_autosave_directory, savedir.lstrip('/')) logger.info('Switched autosave directory to {}, since original location is not writeable.'.format(savedir)) os.makedirs(savedir, exist_ok=True) return savedir
[docs]def get_filename(function, checksum, description, *args): """Get the autosave filename for a specific function call.""" func_desc = function.__name__ for arg in args: if hasattr(arg, '__name__'): func_desc += '_{}'.format(arg.__name__) elif isinstance(arg, functools.partial): func_desc += '_{}'.format(arg.func.__name__) if hasattr(arg, 'frames'): savedir = get_directory(arg.frames) if hasattr(arg, 'description') and arg.description != '': description += '_{}'.format(arg.description) filename = '{}_{}.npz'.format(func_desc.strip('_'), description.strip('_')) return os.path.join(savedir, filename)
[docs]def verify_file(filename, checksum): """Verify if the file matches the function call.""" file_checksum = 0 if os.path.exists(filename): data = np.load(filename) if 'checksum' in data: file_checksum = data['checksum'] return file_checksum == checksum
[docs]def save_data(filename, checksum, data): """Save data and checksum to a file.""" notify('Saving result to file: {}'.format(filename)) try: data = np.array(data) except ValueError: arr = np.empty((len(data),), dtype=object) arr[:] = data data = arr np.savez(filename, checksum=checksum, data=data)
[docs]def load_data(filename): """Load data from a npz file.""" notify('Loading result from file: {}'.format(filename)) fdata = np.load(filename) if 'data' in fdata: return fdata['data'] else: data = tuple(fdata[k] for k in sorted(fdata) if ('arr' in k)) save_data(filename, fdata['checksum'], data) return data
[docs]def autosave_data(nargs, kwargs_keys=None, version=None): """ Enable autosaving of results for a function. Args: nargs: Number of args which are relevant for the calculation. kwargs_keys (opt.): List of keyword arguments which are relevant for the calculation. version (opt.): An optional version number of the decorated function, which replaces the checksum of the function code, hence the checksum does not depend on the function code. """ def decorator_function(function): # make sure too include names of positional arguments in kwargs_keys, # sice otherwise they will be ignored if passed via keyword. # nonlocal kwargs_keys posargs_keys = list(inspect.signature(function).parameters)[:nargs] @functools.wraps(function) def autosave(*args, **kwargs): description = kwargs.pop('description', '') autoload = kwargs.pop('autoload', True) and load_autosave_data if autosave_directory is not None: relevant_args = list(args[:nargs]) if kwargs_keys is not None: for key in [*posargs_keys, *kwargs_keys]: if key in kwargs: relevant_args.append(kwargs[key]) if version is None: csum = legacy_csum = checksum(function, *relevant_args) else: csum = checksum(version, *relevant_args) legacy_csum = checksum(function, *relevant_args) filename = get_filename(function, csum, description, *relevant_args) if autoload and (verify_file(filename, csum) or verify_file(filename, legacy_csum)): result = load_data(filename) else: result = function(*args, **kwargs) save_data(filename, csum, result) else: result = function(*args, **kwargs) return result return autosave return decorator_function