datastore.py

"""Implementation of experiment's data store."""

__all__ = ['DataStoreByVariable', 'DataStoreByScope', 'NameNotFoundInDataStore', 'NameNotFoundInScope', 'EmptyScope',
           'AbstractDataStore']
__author__ = 'Lukas Leufen'
__date__ = '2019-11-22'

import inspect
import types
from abc import ABC
from functools import wraps
from typing import Any, List, Tuple, Dict


class NameNotFoundInDataStore(Exception):
    """Exception that get raised if given name is not found in the entire data store."""

    pass


class NameNotFoundInScope(Exception):
    """Exception that get raised if given name is not found in the provided scope, but can be found in other scopes."""

    pass


class EmptyScope(Exception):
    """Exception that get raised if given scope is not part of the data store."""

    pass


class CorrectScope:
    """
    This class is used as decorator for all class methods, that have scope in parameters.

    After decoration, the scope argument is not required on method call anymore. If no scope parameter is given, this
    decorator automatically adds the default scope=`general` to the arguments. Furthermore, calls like
    `scope=general.sub` are obsolete, because this decorator adds the prefix `general.` if not provided. Therefore, a
    call like `scope=sub` will actually become `scope=general.sub` after passing this decorator.
    """

    def __init__(self, func):
        """Construct decorator."""
        wraps(func)(self)

    def __call__(self, *args, **kwargs):
        """
        Call method of decorator.

        Update tuple if scope argument does not start with `general` or slot `scope=general` into args if not provided
        in neither args nor kwargs.
        """
        f_arg = inspect.getfullargspec(self.__wrapped__)
        pos_scope = f_arg.args.index("scope")
        if len(args) < (len(f_arg.args) - len(f_arg.defaults or "")):
            new_arg = kwargs.pop("scope", "general") or "general"
            args = self.update_tuple(args, new_arg, pos_scope)
        else:
            args = self.update_tuple(args, args[pos_scope], pos_scope, update=True)
        return self.__wrapped__(*args, **kwargs)

    def __get__(self, instance, cls):
        """Create bound method object and supply self argument to the decorated method."""
        return types.MethodType(self, instance)

    @staticmethod
    def correct(arg: str):
        """
        Add leading general prefix.

        :param arg: string argument of scope to add prefix general if necessary

        :return: corrected string
        """
        if not arg.startswith("general"):
            arg = "general." + arg
        return arg

    def update_tuple(self, t: Tuple, new: Any, ind: int, update: bool = False):
        """
        Update single entry n given tuple or slot entry into given position.

        Either update a entry in given tuple t (<old1>, <old2>, <old3>) --(ind=1)--> (<old1>, <new>, <old3>) or slot
        entry into given position (<old1>, <old2>, <old3>) --(ind=1,update=True)--> (<old1>, <new>, <old2>, <old3>). In
        the latter case, length of returned tuple is increased by 1 in comparison to given tuple.

        :param t: tuple to update
        :param new: new element to add to tuple
        :param ind: position to add or slot in
        :param update: updates entry if true, otherwise slot in (default: False)

        :return: updated tuple
        """
        t_new = (*t[:ind], self.correct(new), *t[ind + update:])
        return t_new


class AbstractDataStore(ABC):
    """
    Abstract data store for all settings for the experiment workflow.

    Save experiment parameters for the proceeding run_modules and predefine parameters loaded during the experiment
    setup phase. The data store is hierarchically structured, so that global settings can be overwritten by local
    adjustments.
    """

    def __init__(self):
        """Initialise by creating empty data store."""
        self._store: Dict = {}

    def set(self, name: str, obj: Any, scope: str) -> None:
        """
        Abstract method to add an object to the data store.

        :param name: Name of object to store
        :param obj: The object itself to be stored
        :param scope: the scope / context of the object, under that the object is valid
        """
        pass

    def get(self, name: str, scope: str) -> None:
        """
        Abstract method to get an object from the data store.

        :param name: Name to look for
        :param scope: scope to search the name for
        :return: the stored object
        """
        pass

    def search_name(self, name: str) -> None:
        """
        Abstract method to search for all occurrences of given `name` in the entire data store.

        :param name: Name to look for
        :return: search result
        """
        pass

    def search_scope(self, scope: str) -> None:
        """
        Abstract method to search for all object names that are stored for given scope.

        :param scope: scope to look for
        :return: search result
        """
        pass

    def list_all_scopes(self) -> None:
        """
        Abstract method to list all scopes in data store.

        :return: all found scopes
        """
        pass

    def list_all_names(self) -> None:
        """
        Abstract method to list all names available in the data store.

        :return: all names
        """
        pass

    def clear_data_store(self) -> None:
        """
        Reset entire data store.

        Warning: This will remove all entries of the data store without any exception.
        """
        self._store = {}

    @CorrectScope
    def create_args_dict(self, arg_list: List[str], scope: str) -> Dict:
        """
        Create dictionary from given argument list (as keys) and the stored data inside data store (as values).

        Try to load all stored elements for `arg_list` and create an entry in return dictionary for each valid key
        value pair. Not existing keys from arg_list are skipped. This method works on a single scope only and cannot
        create a dictionary with values from different scopes. Depending on the implementation of the __get__ method,
        all superior scopes are included in the parameter search, if no element is found for the given subscope.

        :param arg_list: list with all elements to look for
        :param scope: the scope to search in

        :return: dictionary with all valid elements from given arg_list as key and the corresponding stored object as
            value.
        """
        args = {}
        for arg in arg_list:
            try:
                args[arg] = self.get(arg, scope)
            except (NameNotFoundInDataStore, NameNotFoundInScope):
                pass
        return args

    @CorrectScope
    def set_from_dict(self, arg_dict: Dict, scope: str) -> None:
        """
        Store multiple objects from dictionary under same `scope`.

        Each object needs to be parsed as key value pair inside the given dictionary. All new entries are stored under
        the same scope.

        :param arg_dict: updates for the data store, provided as key value pairs
        :param scope: scope to store updates
        """
        for (k, v) in arg_dict.items():
            self.set(k, v, scope)


class DataStoreByVariable(AbstractDataStore):
    """
    Data store for all settings for the experiment workflow.

    Save experiment parameters for the proceeding run_modules and predefine parameters loaded during the experiment
    setup phase. The data store is hierarchically structured, so that global settings can be overwritten by local
    adjustments.

    This implementation stores data as

    .. code-block::

        <variable1>
            <scope1>: value
            <scope2>: value
        <variable2>
            <scope1>: value
            <scope3>: value

    """

    @CorrectScope
    def set(self, name: str, obj: Any, scope: str) -> None:
        """
        Store an object `obj` with given `name` under `scope`.

        In the current implementation, existing entries are overwritten.

        :param name: Name of object to store
        :param obj: The object itself to be stored
        :param scope: the scope / context of the object, under that the object is valid
        """
        # open new variable related store with `name` as key if not existing
        if name not in self._store.keys():
            self._store[name] = {}
        self._store[name][scope] = obj

    @CorrectScope
    def get(self, name: str, scope: str) -> Any:
        """
        Retrieve an object with `name` from `scope`.

        If no object can be found in the exact scope, take an iterative look on the levels above. Raise a
        NameNotFoundInDataStore error, if no object with given name can be found in the entire data store. Raise a
        NameNotFoundInScope error, if the object is in the data store but not in the given scope and its levels above
        (could be either included in another scope or a more detailed sub-scope).

        :param name: Name to look for
        :param scope: scope to search the name for

        :return: the stored object
        """
        return self._stride_through_scopes(name, scope)[2]

    @CorrectScope
    def get_default(self, name: str, scope: str, default: Any) -> Any:
        """

        Retrieve an object with `name` from `scope` and return given default if object wasn't found.

        Same functionality like the standard get method. But this method adds a default argument that is returned if no
        data was stored in the data store. Use this function with care, because it will not report any errors and just
        return the given default value. Currently, there is no statement that reports, if the returned value comes from
        the data store or the default value.

        :param name: Name to look for
        :param scope: scope to search the name for
        :param default: default value that is return, if no data was found for given name and scope

        :return: the stored object or the default value
        """
        try:
            return self._stride_through_scopes(name, scope)[2]
        except (NameNotFoundInDataStore, NameNotFoundInScope):
            return default

    @CorrectScope
    def _stride_through_scopes(self, name, scope, depth=0):
        if depth <= scope.count("."):
            local_scope = scope.rsplit(".", maxsplit=depth)[0]
            try:
                return name, local_scope, self._store[name][local_scope]
            except KeyError:
                return self._stride_through_scopes(name, scope, depth + 1)
        else:
            occurrences = self.search_name(name)
            if len(occurrences) == 0:
                raise NameNotFoundInDataStore(f"Couldn't find {name} in data store")
            else:
                raise NameNotFoundInScope(f"Couldn't find {name} in scope {scope} . {name} is only defined in "
                                          f"{occurrences}")

    def search_name(self, name: str) -> List[str]:
        """
        Search for all occurrences of given `name` in the entire data store.

        :param name: Name to look for

        :return: list with all scopes and sub-scopes containing an object stored as `name`
        """
        return sorted(self._store[name] if name in self._store.keys() else [])

    @CorrectScope
    def search_scope(self, scope: str, current_scope_only=True, return_all=False) -> List[str or Tuple]:
        """
        Search for given `scope` and list all object names stored under this scope.

        For an expanded search in all superior scopes, set `current_scope_only=False`. To return the scope and the
        object's value too, set `return_all=True`.

        :param scope: scope to look for
        :param current_scope_only: look only for all names for given scope if true, else search for names from superior
            scopes too.
        :param return_all: return name, definition scope and value if True, else just the name

        :return: list with all object names (if `return_all=False`) or list with tuple of object name, object scope and
            object value ordered by name (if `return_all=True`)
        """
        if current_scope_only:
            names = []
            for (k, v) in self._store.items():
                if scope in v.keys():
                    names.append(k)
            if len(names) > 0:
                if return_all:
                    return sorted([(name, scope, self._store[name][scope]) for name in names], key=lambda tup: tup[0])
                else:
                    return sorted(names)
            else:
                raise EmptyScope(f"Given scope {scope} is not part of the data store. Available scopes are: "
                                 f"{self.list_all_scopes()}")
        else:
            results = []
            for name in self.list_all_names():
                try:
                    res = self._stride_through_scopes(name, scope)
                    if return_all:
                        results.append(res)
                    else:
                        results.append(res[0])
                except (NameNotFoundInDataStore, NameNotFoundInScope):
                    pass
            if return_all:
                return sorted(results, key=lambda tup: tup[0])
            else:
                return sorted(results)

    def list_all_scopes(self) -> List[str]:
        """
        List all available scopes in data store.

        :return: names of all stored objects
        """
        scopes = []
        for v in self._store.values():
            for scope in v.keys():
                if scope not in scopes:
                    scopes.append(scope)
        return sorted(scopes)

    def list_all_names(self) -> List[str]:
        """
        List all names available in the data store.

        :return: all names
        """
        return sorted(self._store.keys())


class DataStoreByScope(AbstractDataStore):
    """
    Data store for all settings for the experiment workflow.

    Save experiment parameters for the proceeding run_modules and predefine parameters loaded during the experiment
    setup phase. The data store is hierarchically structured, so that global settings can be overwritten by local
    adjustments.

    This implementation stores data as

    .. code-block::

        <scope1>
            <variable1>: value
            <variable2>: value
        <scope2>
            <variable1>: value
            <variable3>: value

    """

    @CorrectScope
    def set(self, name: str, obj: Any, scope: str) -> None:
        """
        Store an object `obj` with given `name` under `scope`.

        In the current implementation, existing entries are overwritten.

        :param name: Name of object to store
        :param obj: The object itself to be stored
        :param scope: the scope / context of the object, under that the object is valid
        """
        if scope not in self._store.keys():
            self._store[scope] = {}
        self._store[scope][name] = obj

    @CorrectScope
    def get(self, name: str, scope: str) -> Any:
        """
        Retrieve an object with `name` from `scope`.

        If no object can be found in the exact scope, take an iterative look on the levels above. Raise a
        NameNotFoundInDataStore error, if no object with given name can be found in the entire data store. Raise a
        NameNotFoundInScope error, if the object is in the data store but not in the given scope and its levels above
        (could be either included in another scope or a more detailed sub-scope).

        :param name: Name to look for
        :param scope: scope to search the name for

        :return: the stored object
        """
        return self._stride_through_scopes(name, scope)[2]

    @CorrectScope
    def get_default(self, name: str, scope: str, default: Any) -> Any:
        """
        Retrieve an object with `name` from `scope` and return given default if object wasn't found.

        Same functionality like the standard get method. But this method adds a default argument that is returned if no
        data was stored in the data store. Use this function with care, because it will not report any errors and just
        return the given default value. Currently, there is no statement that reports, if the returned value comes from
        the data store or the default value.

        :param name: Name to look for
        :param scope: scope to search the name for
        :param default: default value that is return, if no data was found for given name and scope

        :return: the stored object or the default value
        """
        try:
            return self._stride_through_scopes(name, scope)[2]
        except (NameNotFoundInDataStore, NameNotFoundInScope):
            return default

    @CorrectScope
    def _stride_through_scopes(self, name, scope, depth=0):
        if depth <= scope.count("."):
            local_scope = scope.rsplit(".", maxsplit=depth)[0]
            try:
                return name, local_scope, self._store[local_scope][name]
            except KeyError:
                return self._stride_through_scopes(name, scope, depth + 1)
        else:
            occurrences = self.search_name(name)
            if len(occurrences) == 0:
                raise NameNotFoundInDataStore(f"Couldn't find {name} in data store")
            else:
                raise NameNotFoundInScope(f"Couldn't find {name} in scope {scope} . {name} is only defined in "
                                          f"{occurrences}")

    def search_name(self, name: str) -> List[str]:
        """
        Search for all occurrences of given `name` in the entire data store.

        :param name: Name to look for

        :return: list with all scopes and sub-scopes containing an object stored as `name`
        """
        keys = []
        for (key, val) in self._store.items():
            if name in val.keys():
                keys.append(key)
        return sorted(keys)

    @CorrectScope
    def search_scope(self, scope: str, current_scope_only: bool = True, return_all: bool = False) -> List[str or Tuple]:
        """
        Search for given `scope` and list all object names stored under this scope.

        For an expanded search in all superior scopes, set `current_scope_only=False`. To return the scope and the
        object's value too, set `return_all=True`.

        :param scope: scope to look for
        :param current_scope_only: look only for all names for given scope if true, else search for names from superior
            scopes too.
        :param return_all: return name, definition scope and value if True, else just the name

        :return: list with all object names (if `return_all=False`) or list with tuple of object name, object scope and
            object value ordered by name (if `return_all=True`)
        """
        if current_scope_only:
            try:
                if return_all:
                    return [(name, scope, self._store[scope][name]) for name in sorted(self._store[scope].keys())]
                else:
                    return sorted(self._store[scope].keys())
            except KeyError:
                raise EmptyScope(f"Given scope {scope} is not part of the data store. Available scopes are: "
                                 f"{self.list_all_scopes()}")
        else:
            results = []
            for name in self.list_all_names():
                try:
                    res = self._stride_through_scopes(name, scope)
                    if return_all:
                        results.append(res)
                    else:
                        results.append(res[0])
                except (NameNotFoundInDataStore, NameNotFoundInScope):
                    pass
            if return_all:
                return sorted(results, key=lambda tup: tup[0])
            else:
                return sorted(results)

    def list_all_scopes(self) -> List[str]:
        """
        List all available scopes in data store.

        :return: names of all stored objects
        """
        return sorted(self._store.keys())

    def list_all_names(self) -> List[str]:
        """
        List all names available in the data store.

        :return: all names
        """
        names = []
        scopes = self.list_all_scopes()
        for scope in scopes:
            for name in self._store[scope].keys():
                if name not in names:
                    names.append(name)
        return sorted(names)