Source code for allensdk.brain_observatory.session_api_utils

import inspect
import warnings

from itertools import zip_longest
from typing import Any, Dict, List

import numpy as np
import pandas as pd


def is_equal(a: Any, b: Any) -> bool:
    """Function to deal with checking if two variables of possibly mixed types
    have the same value."""

    if type(a) != type(b):
        return False

    if isinstance(a, (pd.Series, pd.DataFrame)):
        return a.equals(b)
    elif isinstance(a, np.ndarray):
        return np.array_equal(a, b)
    elif isinstance(a, (list, tuple)):
        for a_elem, b_elem in zip_longest(a, b):
            if not is_equal(a_elem, b_elem):
                return False
        return True
    elif isinstance(a, set):
        for a_elem, b_elem in zip_longest(sorted(a), sorted(b)):
            if not is_equal(a_elem, b_elem):
                return False
        return True
    elif isinstance(a, dict):
        for (a_k, a_v), (b_k, b_v) in zip_longest(sorted(a.items()),
                                                  sorted(b.items())):
            if (a_k != b_k) or (not is_equal(a_v, b_v)):
                return False
        return True
    else:
        return bool(a == b)


class ParamsMixin:
    """This mixin adds parameter management functionality to the class it is
    mixed into.

    This mixin expects that the class it is mixed into will have an __init__
    with type annotated parameters. It also expects for the class to have
    semi-private attributes of the __init__ type annotated parameters.

    Example:

    SomeClassWhereParamManagementIsDesired(ParamsMixin):

        # Managed params should be typed (with simple types if possible)!
        def __init__(self, param_to_ignore, a_param_1: int, a_param_2: float,
                     b_param_1: list):
            # Parameters can be ignored by the mixin
            super().__init__(ignore={'param_to_ignore'})

            # Pay attention to the naming scheme!
            self._a_param_1 = a_param_1
            self._a_param_2 = a_param_2
            self._b_param_1 = b_param_1

        ...

    After being mixed in, methods like 'get_params', 'set_params',
    'needs_data_refresh', and 'clear_updated_params' will be available.
    """

    def __init__(self, ignore: set = {'api'}):
        self._updated_params: set = set()
        self._ignore = ignore

    @classmethod
    def _get_param_signatures(cls) -> List[inspect.Parameter]:
        init = getattr(cls, '__init__')
        if init is object.__init__:
            # Class has a default __init__ and thus no params
            return []
        init_signature = inspect.signature(init)
        # Filter out 'self' and '**kwargs' params
        parameters = [p for p in init_signature.parameters.values()
                      if (p.name != 'self') and (p.kind != p.VAR_KEYWORD)]
        return parameters

    @classmethod
    def _get_param_type_annotations(cls) -> Dict[str, type]:
        parameters = cls._get_param_signatures()
        return {p.name: p.annotation for p in parameters}

    @classmethod
    def _get_param_names(cls) -> List[str]:
        parameters = cls._get_param_signatures()
        return sorted([p.name for p in parameters])

    def get_params(self) -> Dict[str, Any]:
        """Get managed params and their values"""
        out = dict()
        for param in self._get_param_names():
            if param in self._ignore:
                continue
            value = getattr(self, f"_{param}")
            out.update({param: value})
        return out

    def set_params(self, **params):
        """Set managed params"""
        valid_params = self.get_params().keys()
        param_types = self._get_param_type_annotations()
        current_params = self.get_params()

        for param, value in params.items():
            if param in valid_params:
                current_value = current_params[param]

                if isinstance(value, param_types[param]):
                    if not is_equal(current_value, value):
                        setattr(self, f"_{param}", value)
                        self._updated_params.add(param)
                else:
                    warnings.warn(f"The value ({value}) for parameter "
                                  f"'{param}' should be of type "
                                  f"'{param_types[param]}' but is instead "
                                  f"{type(value)}. It will remain as: "
                                  f"{current_value} "
                                  f"({type(current_value)}).",
                                  stacklevel=2)
            else:
                warnings.warn(f"The parameter '{param}' is not valid "
                              f"and is being ignored! "
                              f"Possible params are: {valid_params}",
                              stacklevel=2)

    def needs_data_refresh(self, data_params: set) -> bool:
        """Check if specific params have been updated via `set_params()`"""
        return bool(data_params & self._updated_params)

    def clear_updated_params(self, data_params: set):
        """This method clears 'updated params' whose data have been updated"""
        self._updated_params -= data_params