3
\$\begingroup\$

I wrote a utility library inspired by functools that adds some common operations on functions I frequently use in my projects. As usual I'd appreciate any feed back.

functoolsplus.py

"""More higher-order functions and operations on callable objects."""

from dataclasses import dataclass
from time import perf_counter
from functools import wraps
from sys import exit, stderr    # pylint: disable=W0622
from typing import Any, Callable, IO, Optional, Union


__all__ = [
    'coerce',
    'exiting',
    'exitmethod',
    'instance_of',
    'timeit',
    'wants_instance'
]


Decorator = Callable[[Callable[..., Any]], Callable[..., Any]]


@dataclass
class PerfCounter:
    """Performance counter."""

    start: Optional[float] = None
    end: Optional[float] = None
    on_exit: Optional[Callable[..., Any]] = None

    def __enter__(self):
        self.start = perf_counter()
        return self

    def __exit__(self, typ, value, traceback):
        self.end = perf_counter()

        if self.on_exit is None:
            return None

        if len(self.on_exit.__code__.co_varnames) == 1:
            return self.on_exit(self)

        return self.on_exit()

    @property
    def duration(self) -> float:
        """Return the duration."""
        return self.end - self.start


class exitmethod:   # pylint: disable=C0103
    """Decorator class to create a context manager,
    having the passed function as exit method.
    """

    def __init__(self, function: Callable[..., Any]):
        self.function = function

    def __enter__(self):
        return self

    def __exit__(self, typ, value, traceback):
        if wants_instance(self.function):
            return self.function(self, typ, value, traceback)

        return self.function(typ, value, traceback)


def coerce(typ: type) -> Callable[..., Any]:
    """Converts the return value into the given type."""

    def decorator(function: Callable[..., Any]) -> Callable[..., typ]:
        """Decorates the given function."""
        @wraps(function)
        def wrapper(*args, **kwargs) -> typ:
            """Wraps the respective function."""
            return typ(function(*args, **kwargs))

        wrapper.__annotations__['return'] = typ
        return wrapper

    return decorator


def exiting(function: Callable[..., Any]) -> Callable[..., Any]:
    """Makes a function exit the program with its return code."""

    @wraps(function)
    def wrapper(*args, **kwargs) -> Any:
        """Wraps the respective function."""
        result = function(*args, **kwargs)
        exit(result or 0)

    return wrapper


def instance_of(cls: Union[type, tuple[type]]) -> Callable[[Any], bool]:
    """Returns a callback function to check the instance of an object."""

    return lambda obj: isinstance(obj, cls)


def timeit(file: IO = stderr, flush: bool = False) -> Decorator:
    """Times the execution of the given function."""

    def print_duration(
            function: Callable[..., Any]
    ) -> Callable[[PerfCounter], None]:
        """Prints a perf counter."""

        def inner(ctr: PerfCounter) -> None:
            print('Function', function.__name__, 'took', ctr.duration,
                  file=file, flush=flush)

        return inner

    def decorator(function: Callable[..., Any]) -> Callable[..., Any]:
        """The actual decorator."""
        @wraps(function)
        def wrapper(*args, **kwargs):
            """Wraps the original function."""
            with PerfCounter(on_exit=print_duration(function)):
                return function(*args, **kwargs)

        return wrapper

    return decorator


def wants_instance(function: Callable[..., Any]) -> bool:
    """Determines whether the respective function is considered a method."""

    try:
        return function.__code__.co_varnames[0] == 'self'
    except IndexError:
        return False

Usage examples

The following code are example use cases, to get an idea, how the library utilities are meant to be used.

from time import sleep
from typing import Iterator

import functoolsplus    # Change accordingly


with functoolsplus.PerfCounter() as ctr:
    sleep(3)


print('Performace countner:', ctr, ctr.duration)


@functoolsplus.exitmethod
def value_error_guard(typ, value, traceback):

    if isinstance(value, ValueError):
        print('Doing stuff with ValueError:', value)
        return True


with value_error_guard:
    raise ValueError('Darn it.')


@functoolsplus.coerce(frozenset)
def get_squares(n: int) -> Iterator[int]:
    """Yields square numbers from 0 to n-1."""

    for number in range(n):
        yield number ** 2


input('Press enter to view help(get_squares):')
help(get_squares)
print('Squares:', get_squares(5))


NUMBERS = [1, 2, 3.0, 4, 5, 6.2, '42']
INTEGERS = filter(functoolsplus.instance_of(int), NUMBERS)
print('Numbers:', *NUMBERS)
print('Integers:', *INTEGERS)


@functoolsplus.timeit(flush=True)
def timed_sleep(seconds: float) -> None:

    return sleep(seconds)


timed_sleep(3)


class Foo:

    def foo(self) -> int:
        return 42

    @classmethod
    def bar(cls):
        return cls()

    @staticmethod
    def spamm():
        pass


print('Foo.foo() wants instance:', functoolsplus.wants_instance(Foo.foo))
print('Foo.bar() wants instance:', functoolsplus.wants_instance(Foo.bar))
print('Foo.spamm() wants instance:', functoolsplus.wants_instance(Foo.spamm))


@functoolsplus.exiting
def main(exit_code: int) -> int:

    print('Exiting with exit code:', exit_code)
    return exit_code


main(42)
\$\endgroup\$
4
\$\begingroup\$

I don't find PerfCounter to be well-modelled as a class. start and end aren't valid on construction, which should be a red flag. The current contract also allows for the calling user to initialise start and end themselves, which doesn't make sense; so even if this were to be modelled as a class it shouldn't be a @dataclass. The more natural calling pattern is, I think, a plain function that accepts a function reference, calls it, and returns a float - which already exists, and is called timeit (not your timeit; the built-in timeit). On that note, probably not a good idea to overlap with that name. P.S.: why support a callback? If you're already decorating code, you have access to the part of the code immediately after execution, so why not... just do something there, instead of registering a callback?

Anything that depends on want_instance is a symptom of non-ideal application design. You should already know whether everything is an instance method or not.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.