Module pedantic.decorators.class_decorators

import enum
import types
from dataclasses import is_dataclass
from typing import Callable, Optional, Dict, Type

from pedantic.constants import TYPE_VAR_ATTR_NAME, TYPE_VAR_METHOD_NAME, F, C, TYPE_VAR_SELF
from pedantic.decorators import timer, trace
from pedantic.decorators.fn_deco_pedantic import pedantic, pedantic_require_docstring
from pedantic.env_var_logic import is_enabled
from pedantic.exceptions import PedanticTypeCheckException
from pedantic.type_checking_logic.check_generic_classes import check_instance_of_generic_class_and_get_type_vars, \
    is_instance_of_generic_class


def for_all_methods(decorator: F) -> Callable[[Type[C]], Type[C]]:
    """
        Applies a decorator to all methods of a class.

        Example:

        >>> @for_all_methods(pedantic)
        ... class MyClass(object):
        ...     def m1(self): pass
        ...     def m2(self, x): pass
    """
    def decorate(cls: C) -> C:
        if not is_enabled():
            return cls

        if issubclass(cls, enum.Enum):
            raise PedanticTypeCheckException(f'Enum "{cls}" cannot be decorated with "@pedantic_class". '
                                             f'Enums are not supported yet.')

        if is_dataclass(obj=cls):
            raise PedanticTypeCheckException(f'Dataclass "{cls}" cannot be decorated with "@pedantic_class". '
                                             f'Try to write "@dataclass" over "@pedantic_class".')

        for attr in cls.__dict__:
            attr_value = getattr(cls, attr)

            if isinstance(attr_value, (types.FunctionType, types.MethodType)):
                setattr(cls, attr, decorator(attr_value))
            elif isinstance(attr_value, property):
                prop = attr_value
                wrapped_getter = _get_wrapped(prop=prop.fget, decorator=decorator)
                wrapped_setter = _get_wrapped(prop=prop.fset, decorator=decorator)
                wrapped_deleter = _get_wrapped(prop=prop.fdel, decorator=decorator)
                new_prop = property(fget=wrapped_getter, fset=wrapped_setter, fdel=wrapped_deleter)
                setattr(cls, attr, new_prop)

        _add_type_var_attr_and_method_to_class(cls=cls)
        return cls
    return decorate


def pedantic_class(cls: C) -> C:
    """ Shortcut for @for_all_methods(pedantic) """
    return for_all_methods(decorator=pedantic)(cls=cls)


def pedantic_class_require_docstring(cls: C) -> C:
    """ Shortcut for @for_all_methods(pedantic_require_docstring) """
    return for_all_methods(decorator=pedantic_require_docstring)(cls=cls)


def trace_class(cls: C) -> C:
    """ Shortcut for @for_all_methods(trace) """
    return for_all_methods(decorator=trace)(cls=cls)


def timer_class(cls: C) -> C:
    """ Shortcut for @for_all_methods(timer) """
    return for_all_methods(decorator=timer)(cls=cls)


def _get_wrapped(prop: Optional[F], decorator: F) -> Optional[F]:
    return decorator(prop) if prop is not None else None


def _add_type_var_attr_and_method_to_class(cls: C) -> None:
    def type_vars(self) -> Dict:
        t_vars = {TYPE_VAR_SELF: cls}

        if is_instance_of_generic_class(instance=self):
            type_vars_fifo = getattr(self, TYPE_VAR_ATTR_NAME, dict())
            type_vars_generics = check_instance_of_generic_class_and_get_type_vars(instance=self)
            setattr(self, TYPE_VAR_ATTR_NAME, {**type_vars_fifo, **type_vars_generics, **t_vars})
        else:
            setattr(self, TYPE_VAR_ATTR_NAME, t_vars)

        return getattr(self, TYPE_VAR_ATTR_NAME)

    setattr(cls, TYPE_VAR_METHOD_NAME, type_vars)