defmethod

defmethod.py

# Copyright (c) 2016, 2021 Akshay S <akssri@vakra.xyz> (lispy)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import functools
import types
from typing import Callable, Literal, Optional, Union


def defmethod(
    cl: Union[type, object],
    kind: Optional[Literal['METHOD', 'STATIC', 'CLASS']] = None,
    aroundp=False,
    name: Optional[str] = None,
) -> Callable:
    """decorator to dynamically change class and object methods.

    Parameters
    ----------
    cl: Union[type, Any]
        if `cl` is a class (type), changes class methods, else changes instance methods.
    kind: Optional[str]
        method type in {'method', 'static', 'class'}; default is a 'method'.
        'method': first arg bound to non-type instance (either cl or member of cl).
    aroundp: Optional[False]
        if True, function gets passed previous existing method (if not exist, None) first.
    name: Optional[str]
        name to which function is bound.

    Returns
    -------
    Callable
        closure that takes a function and sets the corresponding method in `cl`

    Example
    -------
    >>> class Foo:
    ...     pass
    ...
    >>> @defmethod(Foo, kind='class')
    ... def say(cl, x):
    ...     print(f'{cl.__name__} says {x}')
    ...
    >>> Foo().say('hah')
    Foo says hah
    >>>
    """
    kind = ('METHOD' if kind is None else kind).upper()

    def add_method(function):
        method_name = name if name is not None else function.__name__
        if aroundp:
            # wrap currently existing method
            prev = getattr(cl, method_name, None)
            function = functools.wraps(prev)(function(prev))
        setattr(
            cl,
            method_name,
            (
                # class: method / static / class
                {'METHOD': lambda x: x, 'STATIC': staticmethod, 'CLASS': classmethod}
                if isinstance(cl, type)
                # instance: method / static / class
                else {
                    'METHOD': lambda f: types.MethodType(f, cl),
                    'STATIC': lambda x: x,
                    'CLASS': lambda f: classmethod(f).__get__(cl),
                }
            ).get(kind)(function),
        )
        return function

    return add_method