Cache.__call__()

Complexity

Conditions

4

Size

Total Lines

30

Duplication

Lines	0
Ratio	0 %

Importance

Changes	1
Bugs	1	Features	0

1 Method

Rating	Name	Duplication	Size	Complexity
A	Cache.inner()	0	12	2

#! /usr/bin/env python
#
# Copyright (C) 2016 Rich Lewis <rl403@cam.ac.uk>
# License: 3-clause BSD

"""
# skchem.features.descriptors.decorators

Decorators for descriptors in scikit-chem.
"""

from functools import wraps

import numpy as np
from rdkit.Chem.rdmolops import GetDistanceMatrix, GetAdjacencyMatrix


def _add_cached_prop(func, prop):

    """ Annotate a function with the cached props. """

    if not hasattr(func, 'caches'):
        func.caches = []
    func.caches.append(prop)


def requires_h_depleted(func):

    """ Decorate a function that requires an h-depleted graph.

    This will check if the molecule argument is h-depleted, and will
    memoize a depleted version if it is not, and pass it to the func.

    Note:
        This decorator should be used first if in combination with dMat etc.
    """

    @wraps(func)
    def inner(mol, *args, **kwargs):

        # if is already h depleted
        if (mol.atoms.atomic_number == 1).sum() == 0:
            return func(mol, *args, **kwargs)

        if not hasattr(mol, '_h_depleted'):
            mol._h_depleted = mol.remove_hs()

        return func(mol._h_depleted, *args, **kwargs)

    _add_cached_prop(inner, '_h_depleted')

    return inner


def requires_h_filled(func):

    """ Decorate a function that requires a h-filled graph.

    This will check if the molecule argument is h-filled, and will
    memoize a filled version if it is not, and pass it to the func.

     Note:
        This decorator should be used first if in combination with dMat etc.

     """

    @wraps(func)
    def inner(mol, *args, **kwargs):
        # if is already h enriched
        if mol.atoms.n_total_hs.sum() == 0:
            return func(mol, *args, **kwargs)

        # if not, memoize the enriched one and pass it
        if not hasattr(mol, '_h_enriched'):
            mol._h_enriched = mol.add_hs()

        return func(mol._h_enriched, *args, **kwargs)

    _add_cached_prop(inner, '_h_enriched')

    return inner


class Cache(object):

    """ Function cache."""

    def __init__(self):
        self.cached = {}

    def __call__(self, func):

        """ Create a decorator to identify a function as returning a cached
        value.

        This can be used for objects  that are nontrivial to generate, but
        are used by many functions.

        Args:
            name (str):
                The name to cache the values under
        """

        name = func.__name__

        @wraps(func)
        def inner(mol, *args, force=False, **kwargs):

invalid syntax

            self.setup_cache(mol)

            # call function if it hasn't already been called.
            if force or name not in mol.cache.keys():
                res = func(mol, *args, **kwargs)
                mol.cache[name] = res

            # return the cached value
            return mol.cache[name]

        self.cached[name] = inner
        return inner

    def retrieve(self, *args_to_inject):

        """ Create a decorator that will inject cached values as arguments.

        Args:
            args (list<str>):
                A list of cachable requirements for this function.
        """

        def outer(func):

            @wraps(func)
            def inner(mol, *args, **kwargs):

                self.setup_cache(mol)

                # look up cached values, or produce them if not.
                # inject the cached values
                args = tuple(mol.cache.get(arg.__name__,
                                           self.cached[arg.__name__](mol))
                             for arg in args_to_inject) + args

                return func(mol, *args, **kwargs)
            return inner
        return outer

    @staticmethod
    def setup_cache(mol):
        if not hasattr(mol, 'cache'):
            mol.cache = {}

    @staticmethod
    def teardown_cache(mol):
        if hasattr(mol, 'cache'):
            del mol.cache

cache = Cache()


def requires_dmat(func):
    """ Decorate a function that requires a distance matrix. """
    @wraps(func)
    def inner(mol, *args, **kwargs):
        if not hasattr(mol, '_dMat'):
            dmat = GetDistanceMatrix(mol)
            dmat[dmat > 100] = np.nan
            mol._dMat = dmat
        return func(mol, *args, **kwargs)

    _add_cached_prop(inner, '_dMat')

    return inner


def requires_amat(func):
    """ Decorate a function that requires an adjacency matrix. """
    @wraps(func)
    def inner(mol, *args, **kwargs):
        if not hasattr(mol, '_adjMat'):
            mol._adjMat = GetAdjacencyMatrix(mol)
        return func(mol, *args, **kwargs)

    _add_cached_prop(inner, 'amat')

    return inner


def requires_bo_amat(func):
    """ Decorate a function thate requires a bond order adjacency, matrix. """
    @wraps(func)
    def inner(mol, *args, **kwargs):
        if not hasattr(mol, '_bo_amat'):
            mol._bo_amat = GetAdjacencyMatrix(mol, useBO=1)
        return func(mol, *args, **kwargs)

    _add_cached_prop(inner, '_bo_amat')

    return inner


def requires_degrees(func):
    """ Decorate a function that requires a degree vector. """
    @wraps(func)
    def inner(mol, *args, **kwargs):
        if not hasattr(mol, '_degrees'):
            mol._degrees = np.array([atom.degree for atom in mol.atoms])
        return func(mol, *args, **kwargs)

    _add_cached_prop(inner, '_degrees')

    return inner